# SynJax: Structured Probability Distributions for JAX

arXiv: https://arxiv.org/abs/2308.03291 · HTML: https://mlsanity.com/paper/2308.03291/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 87 references that are in the map.
- First posted on arXiv 7 Aug 2023.

## Builds on (most cited first, 30 of 37)

- [End-to-end Sequence Labeling via Bi-directional LSTM-CNNs-CRF](https://mlsanity.com/paper/1603.01354/) (2016)
- [Soft-DTW: a Differentiable Loss Function for Time-Series](https://mlsanity.com/paper/1703.01541/) (2017)
- [Differentiation of Blackbox Combinatorial Solvers](https://mlsanity.com/paper/1912.02175/) (2019)
- [Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro](https://mlsanity.com/paper/1912.11554/) (2019)
- [Learning Latent Permutations with Gumbel-Sinkhorn Networks](https://mlsanity.com/paper/1802.08665/) (2018)
- [Equinox: neural networks in JAX via callable PyTrees and filtered transformations](https://mlsanity.com/paper/2111.00254/) (2021)
- [Interpretable Neural Predictions with Differentiable Binary Variables](https://mlsanity.com/paper/1905.08160/) (2019)
- [Multilingual Constituency Parsing with Self-Attention and Pre-Training](https://mlsanity.com/paper/1812.11760/) (2018)
- [Recurrent Neural Network Grammars](https://mlsanity.com/paper/1602.07776/) (2016)
- [TensorFlow Distributions](https://mlsanity.com/paper/1711.10604/) (2017)
- [Einsum Networks: Fast and Scalable Learning of Tractable Probabilistic Circuits](https://mlsanity.com/paper/2004.06231/) (2020)
- [Implicit MLE: Backpropagating Through Discrete Exponential Family Distributions](https://mlsanity.com/paper/2106.01798/) (2021)
- [Structured Attention Networks](https://mlsanity.com/paper/1702.00887/) (2017)
- [A Differential Approach to Inference in Bayesian Networks](https://mlsanity.com/paper/1301.3847/) (2013)
- [Compound Probabilistic Context-Free Grammars for Grammar Induction](https://mlsanity.com/paper/1906.10225/) (2019)
- [Graph Convolutional Encoders for Syntax-aware Neural Machine Translation](https://mlsanity.com/paper/1704.04675/) (2017)
- [Gradient Estimation with Stochastic Softmax Tricks](https://mlsanity.com/paper/2006.08063/) (2020)
- [A Minimal Span-Based Neural Constituency Parser](https://mlsanity.com/paper/1705.03919/) (2017)
- [Differentiable Dynamic Programming for Structured Prediction and Attention](https://mlsanity.com/paper/1802.03676/) (2018)
- [Transformer Grammars: Augmenting Transformer Language Models with Syntactic Inductive Biases at Scale](https://mlsanity.com/paper/2203.00633/) (2022)
- [Torch-Struct: Deep Structured Prediction Library](https://mlsanity.com/paper/2002.00876/) (2020)
- [AMR Parsing as Graph Prediction with Latent Alignment](https://mlsanity.com/paper/1805.05286/) (2018)
- [Temporal Parallelization of Bayesian Smoothers](https://mlsanity.com/paper/1905.13002/) (2019)
- [Cooperative Learning of Disjoint Syntax and Semantics](https://mlsanity.com/paper/1902.09393/) (2019)
- [Unsupervised Neural Hidden Markov Models](https://mlsanity.com/paper/1609.09007/) (2016)
- [Differentiable Perturb-and-Parse: Semi-Supervised Parsing with a Structured Variational Autoencoder](https://mlsanity.com/paper/1807.09875/) (2018)
- [Predicting Target Language CCG Supertags Improves Neural Machine Translation](https://mlsanity.com/paper/1702.01147/) (2017)
- [Efficient Computation of Expectations under Spanning Tree Distributions](https://mlsanity.com/paper/2008.12988/) (2020)
- [Please Mind the Root: Decoding Arborescences for Dependency Parsing](https://mlsanity.com/paper/2010.02550/) (2020)
- [LinearFold: linear-time approximate RNA folding by 5'-to-3' dynamic programming and beam search](https://mlsanity.com/paper/2001.04020/) (2019)

## Built on by (most cited first, 2 of 2)

- [Prefix Parsing is Just Parsing](https://mlsanity.com/paper/2604.21191/) (2026)
- [Syntactic Belief Update as the Driver of Garden Path Processing Difficulty](https://mlsanity.com/paper/2606.27206/) (2026)

## Abstract (arXiv)

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. This is done by exploiting the connection between algorithms for automatic differentiation and probabilistic inference. With SynJax we can build large-scale differentiable models that explicitly model structure in the data. The code is available at https://github.com/google-deepmind/synjax