title: SynJax: Structured Probability Distributions for JAX
summary: The development of deep learning software libraries enabled significant
  progress in the field by allowing users to focus on modeling, while letting the
  library to take care of the tedious and time-consuming task of optimizing
  execution for modern hardware accelerators. However, this has benefited only
  particular types of deep learning models, such as Transformers, whose
  primitives map easily to the vectorized computation. The models that explicitly
  account for structured objects, such as trees and segmentations, did not
  benefit equally because they require custom algorithms that are difficult to
  implement in a vectorized form.
    SynJax directly addresses this problem by providing an efficient vectorized
  implementation of inference algorithms for structured distributions covering
  alignment, tagging, segmentation, constituency trees and spanning trees. With
  SynJax we can build large-scale differentiable models that explicitly model
  structure in the data. The code is available at
  https://github.com/deepmind/synjax.
paper_authors._id: 64d1d93a749587dbe0665961
paper_authors.name: Miloš Stanojević
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2023-08-09T00:42:48.379Z
paper_authors._id: 64d2626b07abea8697becb14
paper_authors.avatarUrl: /avatars/475fb745858efda6eb6b54c60d6ac79e.svg
paper_authors.fullname: Milos Stanojevic
paper_authors.type: user
paper_authors.user: stanojevic
paper_authors._id: 64d1d93a749587dbe0665962
paper_authors.name: Laurent Sartran
paper_title: SynJax: Structured Probability Distributions for JAX
paper_summary: The development of deep learning software libraries enabled significant
  progress in the field by allowing users to focus on modeling, while letting the
  library to take care of the tedious and time-consuming task of optimizing
  execution for modern hardware accelerators. However, this has benefited only
  particular types of deep learning models, such as Transformers, whose
  primitives map easily to the vectorized computation. The models that explicitly
  account for structured objects, such as trees and segmentations, did not
  benefit equally because they require custom algorithms that are difficult to
  implement in a vectorized form.
    SynJax directly addresses this problem by providing an efficient vectorized
  implementation of inference algorithms for structured distributions covering
  alignment, tagging, segmentation, constituency trees and spanning trees. With
  SynJax we can build large-scale differentiable models that explicitly model
  structure in the data. The code is available at
  https://github.com/deepmind/synjax.
paper_ai_keywords: Transformers vectorized computation structured objects trees segmentations custom algorithms inference algorithms structured distributions alignment tagging Segmentation large-scale differentiable models