torchgfn: A PyTorch GFlowNet library
Abstract
The growing popularity of generative flow networks (GFlowNets or GFNs) from a range of researchers with diverse backgrounds and areas of expertise necessitates a library that facilitates the testing of new features (e.g., training losses and training policies) against standard benchmark implementations, or on a set of common environments. We present torchgfn, a PyTorch library that aims to address this need. Its core contribution is a modular and decoupled architecture which treats environments, neural network modules, and training objectives as interchangeable components. This provides users with a simple yet powerful API to facilitate rapid prototyping and novel research. Multiple examples are provided, replicating and unifying published results. The library is focused on maintainability and reliability, with static type checking, style enforcement, and tests ensuring library integrity during ongoing development using continuous integration. The library is available on GitHub (https://anonymous.4open.science/r/torchgfn-anon-1D7F/) and on pypi (link hidden).