# gfnx: Fast and Scalable Library for Generative Flow Networks in JAX

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 42 of its 70 references that are in the map.
- First posted on arXiv 20 Nov 2025.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation](https://mlsanity.com/paper/2106.04399/) (2021)
- [GFlowNet Foundations](https://mlsanity.com/paper/2111.09266/) (2021)
- [Trajectory balance: Improved credit assignment in GFlowNets](https://mlsanity.com/paper/2201.13259/) (2022)
- [Learning Bayesian Networks: The Combination of Knowledge and Statistical Data](https://mlsanity.com/paper/1302.6815/) (2013)
- [Biological Sequence Design with GFlowNets](https://mlsanity.com/paper/2203.04115/) (2022)
- [Bayesian Structure Learning with Generative Flow Networks](https://mlsanity.com/paper/2202.13903/) (2022)
- [Equinox: neural networks in JAX via callable PyTrees and filtered transformations](https://mlsanity.com/paper/2111.00254/) (2021)
- [Learning GFlowNets from partial episodes for improved convergence and stability](https://mlsanity.com/paper/2209.12782/) (2022)
- [Generative Flow Networks for Discrete Probabilistic Modeling](https://mlsanity.com/paper/2202.01361/) (2022)
- [GFlowNets and variational inference](https://mlsanity.com/paper/2210.00580/) (2022)
- [Discovered Policy Optimisation](https://mlsanity.com/paper/2210.05639/) (2022)
- [Multi-Objective GFlowNets](https://mlsanity.com/paper/2210.12765/) (2022)
- [Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning](https://mlsanity.com/paper/2402.16801/) (2024)
- [Better Training of GFlowNets with Local Credit and Incomplete Trajectories](https://mlsanity.com/paper/2302.01687/) (2023)
- [Towards Understanding and Improving GFlowNet Training](https://mlsanity.com/paper/2305.07170/) (2023)
- [Generative Flow Networks as Entropy-Regularized RL](https://mlsanity.com/paper/2310.12934/) (2023)
- [JaxMARL: Multi-Agent RL Environments and Algorithms in JAX](https://mlsanity.com/paper/2311.10090/) (2023)
- [Learning Gaussian Networks](https://mlsanity.com/paper/1302.6808/) (2013)
- [How Many Random Seeds? Statistical Power Analysis in Deep Reinforcement Learning Experiments](https://mlsanity.com/paper/1806.08295/) (2018)
- [Hyperparameters in Reinforcement Learning and How To Tune Them](https://mlsanity.com/paper/2306.01324/) (2023)
- [XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX](https://mlsanity.com/paper/2312.12044/) (2023)
- [Pgx: Hardware-Accelerated Parallel Game Simulators for Reinforcement Learning](https://mlsanity.com/paper/2303.17503/) (2023)
- [Local Search GFlowNets](https://mlsanity.com/paper/2310.02710/) (2023)
- [Molecular Mechanics-Driven Graph Neural Network with Multiplex Graph for Molecular Structures](https://mlsanity.com/paper/2011.07457/) (2020)
- [Discrete Probabilistic Inference as Control in Multi-path Environments](https://mlsanity.com/paper/2402.10309/) (2024)
- [RGFN: Synthesizable Molecular Generation Using GFlowNets](https://mlsanity.com/paper/2406.08506/) (2024)
- [Bayesian learning of Causal Structure and Mechanisms with GFlowNets and Variational Bayes](https://mlsanity.com/paper/2211.02763/) (2022)
- [Maximum entropy GFlowNets with soft Q-learning](https://mlsanity.com/paper/2312.14331/) (2023)

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

- [Learning Shortest Paths with Generative Flow Networks](https://mlsanity.com/paper/2603.01786/) (2026)
- [Proximal Policy Optimization for Amortized Discrete Sampling](https://mlsanity.com/paper/2606.15793/) (2026)
- [Learning a Mixture of GFlowNets](https://mlsanity.com/paper/2610.07562/) (2026)

## Abstract (arXiv)

In this paper, we present gfnx, a fast and scalable package for training and evaluating Generative Flow Networks (GFlowNets) written in JAX. gfnx provides an extensive set of environments and metrics for benchmarking, accompanied with single-file implementations of core objectives for training GFlowNets. We include synthetic hypergrids, multiple sequence generation environments with various editing regimes and particular reward designs for molecular generation, phylogenetic tree construction, Bayesian structure learning, and sampling from the Ising model energy. Across different tasks, gfnx achieves significant wall-clock speedups compared to Pytorch-based benchmarks (such as torchgfn library) and author implementations. For example, gfnx achieves up to 55 times speedup on CPU-based sequence generation environments, and up to 80 times speedup with the GPU-based Bayesian network structure learning setup. Our package provides a diverse set of benchmarks and aims to standardize empirical evaluation and accelerate research and applications of GFlowNets. The library is available on GitHub (https://github.com/d-tiapkin/gfnx) and on pypi (https://pypi.org/project/gfnx/). Documentation is available on https://gfnx.readthedocs.io.