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.

翻译：本文介绍了gfnx——一个基于JAX编写的、用于训练和评估生成流网络（GFlowNets）的快速可扩展软件包。gfnx提供了丰富的基准测试环境和评估指标，并采用单文件形式实现了训练GFlowNets的核心目标函数。该库包含合成超网格环境、多种序列生成环境（涵盖不同编辑机制与特定奖励设计），以及面向分子生成、系统发育树构建、贝叶斯结构学习和伊辛模型能量采样的专用环境。在不同任务中，相较于基于PyTorch的基准实现（如torchgfn库）及作者原始实现，gfnx均取得显著的实际运行加速效果。例如，在基于CPU的序列生成环境中实现高达55倍加速，在基于GPU的贝叶斯网络结构学习场景中实现高达80倍加速。本软件包提供多样化的基准测试集，旨在规范实证评估流程，推动GFlowNets的研究与应用发展。该库已发布于GitHub（https://github.com/d-tiapkin/gfnx）和PyPI（https://pypi.org/project/gfnx/），技术文档详见https://gfnx.readthedocs.io。