# NetKet 3: Machine Learning Toolbox for Many-Body Quantum Systems

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

## Key facts

- Cited by 28 papers in the map (data through 14 Oct 2026).
- Builds on 11 of its 99 references that are in the map.
- First posted on arXiv 20 Dec 2021.

## Builds on (most cited first, 11 of 11)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [This Looks Like That: Deep Learning for Interpretable Image Recognition](https://mlsanity.com/paper/1806.10574/) (2018)
- [Quantum Natural Gradient](https://mlsanity.com/paper/1909.02108/) (2019)
- [Solving Statistical Mechanics Using Variational Autoregressive Networks](https://mlsanity.com/paper/1809.10606/) (2018)
- [Deep autoregressive models for the efficient variational simulation of many-body quantum systems](https://mlsanity.com/paper/1902.04057/) (2019)
- [Neural-Network Quantum States, String-Bond States, and Chiral Topological States](https://mlsanity.com/paper/1710.04045/) (2017)
- [Neural network wave functions and the sign problem](https://mlsanity.com/paper/2002.04613/) (2020)
- [Geometry of learning neural quantum states](https://mlsanity.com/paper/1910.11163/) (2019)
- [Summarizing CPU and GPU Design Trends with Product Data](https://mlsanity.com/paper/1911.11313/) (2019)

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

- [Wasserstein Quantum Monte Carlo: A Novel Approach for Solving the Quantum Many-Body Schr"odinger Equation](https://mlsanity.com/paper/2307.07050/) (2023)
- [ANTN: Bridging Autoregressive Neural Networks and Tensor Networks for Quantum Many-Body Simulation](https://mlsanity.com/paper/2304.01996/) (2023)
- [m^ast of two-dimensional electron gas: a neural canonical transformation study](https://mlsanity.com/paper/2201.03156/) (2022)
- [Ab-initio variational wave functions for the time-dependent many-electron Schr"odinger equation](https://mlsanity.com/paper/2403.07447/) (2024)
- [Forward Laplacian: A New Computational Framework for Neural Network-based Variational Monte Carlo](https://mlsanity.com/paper/2307.08214/) (2023)
- [Are queries and keys always relevant? A case study on Transformer wave functions](https://mlsanity.com/paper/2405.18874/) (2024)
- [When can classical neural networks represent quantum states?](https://mlsanity.com/paper/2410.23152/) (2024)
- [Interpretable representation learning of quantum data enabled by probabilistic variational autoencoders](https://mlsanity.com/paper/2506.11982/) (2025)
- [Approximately-symmetric neural networks for quantum spin liquids](https://mlsanity.com/paper/2405.17541/) (2024)
- [Q-Flow: Generative Modeling for Differential Equations of Open Quantum Dynamics with Normalizing Flows](https://mlsanity.com/paper/2302.12235/) (2023)
- [Empirical Sample Complexity of Neural Network Mixed State Reconstruction](https://mlsanity.com/paper/2307.01840/) (2023)
- [Deep Stochastic Mechanics](https://mlsanity.com/paper/2305.19685/) (2023)
- [Functional Neural Wavefunction Optimization](https://mlsanity.com/paper/2507.10835/) (2025)
- [Discovering quantum phenomena with Interpretable Machine Learning](https://mlsanity.com/paper/2604.16015/) (2026)
- [Learning ground states of gapped quantum Hamiltonians with Kernel Methods](https://mlsanity.com/paper/2303.08902/) (2023)
- [Probabilistic Computers for Neural Quantum States](https://mlsanity.com/paper/2512.24558/) (2025)
- [Stochastic Reconfiguration as Statistical Filtering for Overparameterized Neural Quantum States](https://mlsanity.com/paper/2609.23334/) (2026)
- [An Empirical Study of Quantum Dynamics as a Ground State Problem with Neural Quantum States](https://mlsanity.com/paper/2206.09241/) (2022)
- [Neural network quantum state with proximal optimization: a ground-state searching scheme based on variational Monte Carlo](https://mlsanity.com/paper/2210.16493/) (2022)
- [Hybrid Ground-State Quantum Algorithms based on Neural Schr"odinger Forging](https://mlsanity.com/paper/2307.02633/) (2023)
- [Variational decision diagrams for quantum-inspired machine learning applications](https://mlsanity.com/paper/2502.04271/) (2025)
- [DeepQuark: A Deep-Neural-Network Approach to Multiquark Bound States](https://mlsanity.com/paper/2506.20555/) (2025)
- [Fermionic neural Gibbs states](https://mlsanity.com/paper/2512.04663/) (2025)
- [Exploring the Effect of Basis Rotation on NQS Performance](https://mlsanity.com/paper/2512.17893/) (2025)
- [Tensor-Augmented Convolutional Neural Networks: Enhancing Expressivity with Generic Tensor Kernels](https://mlsanity.com/paper/2604.08072/) (2026)
- [Hyperbolic Restricted Boltzmann Machine Neural Quantum State](https://mlsanity.com/paper/2609.26032/) (2026)
- [Apparent Compression, Real Stability: The Intrinsic Dimension of Learning a Quantum Wavefunction](https://mlsanity.com/paper/2609.33193/) (2026)
- [Probabilistic Algorithms for Ising Machines from Optimization to Generative AI](https://mlsanity.com/paper/2610.03972/) (2026)

## Abstract (arXiv)

We introduce version 3 of NetKet, the machine learning toolbox for many-body quantum physics. NetKet is built around neural-network quantum states and provides efficient algorithms for their evaluation and optimization. This new version is built on top of JAX, a differentiable programming and accelerated linear algebra framework for the Python programming language. The most significant new feature is the possibility to define arbitrary neural network ans"atze in pure Python code using the concise notation of machine-learning frameworks, which allows for just-in-time compilation as well as the implicit generation of gradients thanks to automatic differentiation. NetKet 3 also comes with support for GPU and TPU accelerators, advanced support for discrete symmetry groups, chunking to scale up to thousands of degrees of freedom, drivers for quantum dynamics applications, and improved modularity, allowing users to use only parts of the toolbox as a foundation for their own code.
