mlsanityData to 7 Oct 2026

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

arXiv 2511.16592 ↗First posted 20 Nov 2025cs.LGcited by 2

Not on the map. Why?

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.

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

Builds on

42 of 70 references.

References · 42

  1. 2017Attention Is All You Needcited by 59,932
  2. 2021Flow Network based Generative Models for Non-Iterative Diverse Candidate Generationcited by 384
  3. 2021Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulationcited by 375
  4. 2021GFlowNet Foundationscited by 286
  5. 2022Trajectory balance: Improved credit assignment in GFlowNetscited by 244
  6. 2013Learning Bayesian Networks: The Combination of Knowledge and Statistical Datacited by 185
  7. 2022Biological Sequence Design with GFlowNetscited by 172
  8. 2022Bayesian Structure Learning with Generative Flow Networkscited by 153
  9. 2021Equinox: neural networks in JAX via callable PyTrees and filtered transformationscited by 143
  10. 2022Learning GFlowNets from partial episodes for improved convergence and stabilitycited by 129
  11. 2022Generative Flow Networks for Discrete Probabilistic Modelingcited by 113
  12. 2022GFlowNets and variational inferencecited by 103
30 more
  1. 2022Discovered Policy Optimisationcited by 103
  2. 2022Multi-Objective GFlowNetscited by 100
  3. 2023Better Training of GFlowNets with Local Credit and Incomplete Trajectoriescited by 83
  4. 2024Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learningcited by 80
  5. 2023Towards Understanding and Improving GFlowNet Trainingcited by 74
  6. 2013Learning Gaussian Networkscited by 62
  7. 2023Generative Flow Networks as Entropy-Regularized RLcited by 60
  8. 2023JaxMARL: Multi-Agent RL Environments and Algorithms in JAXcited by 59
  9. 2023Hyperparameters in Reinforcement Learning and How To Tune Themcited by 58
  10. 2018How Many Random Seeds? Statistical Power Analysis in Deep Reinforcement Learning Experimentscited by 56
  11. 2023XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAXcited by 54
  12. 2023Pgx: Hardware-Accelerated Parallel Game Simulators for Reinforcement Learningcited by 45
  13. 2023Local Search GFlowNetscited by 45
  14. 2020Molecular Mechanics-Driven Graph Neural Network with Multiplex Graph for Molecular Structurescited by 44
  15. 2024Discrete Probabilistic Inference as Control in Multi-path Environmentscited by 40
  16. 2024RGFN: Synthesizable Molecular Generation Using GFlowNetscited by 39
  17. 2022Bayesian learning of Causal Structure and Mechanisms with GFlowNets and Variational Bayescited by 30
  18. 2023Maximum entropy GFlowNets with soft Q-learningcited by 30
  19. 2022EvoJAX: Hardware-Accelerated Neuroevolutioncited by 29
  20. 2024Adaptive teachers for amortized samplerscited by 23
  21. 2024Discovering Temporally-Aware Reinforcement Learning Algorithmscited by 18
  22. 2024Position: Benchmarking is Limited in Reinforcement Learning Researchcited by 15
  23. 2023Order-Preserving GFlowNetscited by 13
  24. 2023A Theory of Non-Acyclic Generative Flow Networkscited by 12
  25. 2024Optimizing Backward Policies in GFlowNets via Trajectory Likelihood Maximizationcited by 12
  26. 2024JaxUED: A simple and useable UED library in Jaxcited by 10
  27. 2023torchgfn: A PyTorch GFlowNet librarycited by 9
  28. 2024Improving GFlowNets with Monte Carlo Tree Searchcited by 8
  29. 2025Revisiting Non-Acyclic GFlowNets in Discrete Environmentscited by 8
  30. 2024Random Policy Evaluation Uncovers Policies of Generative Flow Networkscited by 5

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2026: 2.

2 papers.

  1. 2026Learning Shortest Paths with Generative Flow Networkscited by 2
  2. 2026Proximal Policy Optimization for Amortized Discrete Samplingcited by 0