# TORAX turns a tokamak core transport model into a differentiable JAX program

> TORAX is a one dimensional tokamak core transport simulator built on JAX, aimed at forward modelling, pulse design, trajectory optimisation and controller work. Its differentiability and its coupling to neural network surrogates are the reason to use it, and the pinned scientific dependencies are the reason to think twice.

**google-deepmind/torax** — TORAX: Tokamak transport simulation in JAX

- Repository: https://github.com/google-deepmind/torax
- Website: https://torax.readthedocs.io
- Stars: 744 · Forks: 154
- Language: Python
- License: NOASSERTION
- Published: 2026-10-21 · Updated: 2026-10-21 · Language: en
- Canonical page: https://hysenlabs.com/projects/google-deepmind-torax

## A core transport code built so that gradients flow through it

TORAX simulates transport in the core of a tokamak, which is the hard part of fusion modelling: the region where confinement and current profiles decide whether a discharge holds. The project describes itself as a differentiable tokamak core transport simulator aimed at fast and accurate forward modelling, pulse design, trajectory optimisation and controller design. The differentiability is the point rather than a side effect, because a gradient-based solver needs to treat the whole time evolution as something a computer can differentiate end to end.
The README gives the reason in concrete terms. Auto-differentiation means you do not have to hand-derive Jacobians when you add a physics model, an ML surrogate or a new parameter input, and the same machinery gives you sensitivity of the simulation output to arbitrary inputs. That combination of gradient access and Python is what makes the project a natural host for data-driven parameter identification of semi-empirical models.
It also makes clear who this is for. The stated motivations are open source and extensibility, flexible coupling into existing workflows, and serving as a natural framework for coupling machine learned surrogates of physics models. The documentation is hosted on readthedocs, and the project is not an officially supported Google product.

## Coupled transport equations on a finite volume grid

The physics core couples four things: ion and electron heat transport, electron particle transport, and current diffusion. These are solved as coupled partial differential equations, discretised with the finite volume method. Poloidal flux boundary conditions can be posed either on total current or on loop voltage at the last closed flux surface, which is what lets the same solver handle current-driven and voltage-driven scenarios.
Source terms complete the energy balance: ohmic power, ion-electron heat exchange, fusion power, Bremsstrahlung, impurity line radiation, and an ICRH machine learning surrogate whose coverage the README qualifies as covering limited regimes so far. Neoclassical bootstrap current and conductivity come from the analytical Sauter model, which is a closed form treatment rather than a neural network.
Two more features handle the edge of the discharge. Simple pedestal models set internal boundary conditions through a local adaptive source, and sawtooth triggering with profile redistribution is modelled. On top of that, additional heating and current drive sources can come from analytical models or from prescribed data that you supply, so the code does not lock you into one source set.

## Three solver strategies selected from configuration

The solver choice is a configuration decision, and the README lists three options. A linear solver with Pereverzev-Corrigan terms and a predictor-corrector method handles the cases where a linear solve is close enough. A nonlinear solver using Newton-Raphson handles the harder cases. A third option runs the nonlinear problem through optimisation using the jaxopt library, which turns each timestep into an optimisation rather than a root find.
The predictor-corrector and linear options exist because full Newton iterations at every timestep cost more than many simulations can afford. The README explains the progress flag in terms of this: the log_progress flag outputs simulation time, dt, and the number of solver iterations carried out at each timestep, including dt backtracking when the nonlinear solver has to shrink a step to converge. That detail matters more than it looks, because a silent fallback to smaller timesteps changes the cost of a run without changing the answer.
Error checking is a related knob, and the documentation ties it to the compilation cache. The TORAX_ERRORS_ENABLED environment variable turns on error checking in internal routines and is used for debugging, but its default is false since it is incompatible with the persistent compilation cache. The same trade-off appears in JAX_DISABLE_JIT, which stops JAX from compiling internal TORAX functions and is also defaulted to false, and which exists for debugging rather than for production.

## Installing TORAX from PyPI and running your first case

TORAX needs Python 3.12 or greater, and the README asks you to install tkinter before anything else, since some parts of the tooling rely on it:

```shell
sudo apt-get install python3-tk
```

The project recommends working inside a virtual environment. The README lists the upgrade of pip and the installation of virtualenv as separate steps:

```shell
pip install --upgrade pip
```

```shell
pip install virtualenv
```

After that you create and activate an environment for the project:

```shell
python3 -m venv toraxvenv
source toraxvenv/bin/activate
```

The simplest route to a working simulator is PyPI, and the smoke test the README suggests runs the basic configuration and exits straight away, so you get an exit status rather than a finished simulation:

```shell
pip install torax
```

```shell
run_torax --config=examples/basic_config.py --quit
```

A full run uses the same default configuration file without the quit flag, and a terminal progress bar reports the current time and the percentage completed:

```shell
run_torax --config='examples/basic_config.py'
```

## Geometry arrives as CHEASE, FBT or EQDSK files

TORAX is a one dimensional code, so it needs the equilibrium geometry from somewhere else. General geometry is provided through equilibrium files in CHEASE, FBT or EQDSK format, and the reading of those formats is handled by the eqdsk package listed among the dependencies. Time dependent geometry is supported by handing the code a time series of geometry files, which is how a scenario that changes shape during the discharge is expressed.
For testing and demonstration, the repository ships a single CHEASE equilibrium file in the data/third_party/geo directory. The README describes it as an ITER hybrid scenario equilibrium based on simulations from a published paper, and states that it was obtained from PINT, with a PINT licence file in the same directory. That is a limited but workable starting point: one scenario to check that your installation behaves, not a library of geometries to explore.
Output handling follows the same pattern. The run_torax entry point writes a time history of the simulation state, and the manifest also installs a second console script, plot_torax, for plotting runs. Output format is handled by the xarray, h5py, h5netcdf and netcdf4 dependencies, and the configuration exposes an output_dir flag that overrides the default output directory.

## Where TORAX stops short, and what it is checked against

The honest limit is the ICRH surrogate. The README itself qualifies its coverage as limited regimes, which means results that depend on that source outside the trained range carry the accuracy of the surrogate rather than the accuracy of a physical model. Anyone using TORAX for ICRH-dominated scenarios needs to know which regimes the surrogate covers.
The second limit is that a one dimensional core model is not a whole-device model. TORAX handles the core, and the sawtooth and pedestal features are simple local models. Anything that depends on edge localised modes, disruptions or global MHD stability is outside what this code is built to answer, and no part of the documentation claims otherwise.
The alternative to compare against is RAPTOR, the tokamak transport simulator that the README says the model implementation was verified against through direct comparison of simulation outputs. That is the difference in approach worth understanding: TORAX is not merely a faster RAPTOR, it is a transport code whose equations are written in JAX so that the same model can be differentiated, and whose verification path is agreement with an existing simulator on physics rather than agreement on gradients.
The same substitution of machine learning for physics appears elsewhere. Turbulent transport can come from the QuaLiKiz neural network surrogates, both QLKNN_7_11 and QLKNN10D, provided through the fusion_surrogates package. That gives you a fast surrogate model, and it moves the modelling uncertainty into the surrogate rather than removing it.

## A pinned scientific stack is the upgrade cost

The packaging tells you how heavy this project is. pyproject.toml pulls in jax and jaxlib, jaxopt, flax, equinox, chex, jaxtyping, scipy, numpy, xarray, matplotlib, plotly, h5py, h5netcdf, netcdf4, eqdsk, pydantic, tqdm and imas-python, plus a pinned build of fusion_surrogates. Two of those pins are hard: fusion_surrogates is locked to version 0.4.7 and imas-python to version 2.1.0, and h5py carries an upper limit of 3.15.0 with a comment saying the limit should be removed once the bug behind it is fixed.

That last comment is worth reading twice, because it means a dependency constraint exists purely to work around an unresolved defect. An environment that already pins h5py above that range will conflict with TORAX, and the fix is not yours to make.

The dependency tree is also a maintenance cost. Python 3.12 is the floor, and every package above moves on its own schedule. The README points contributors at a contributing page and a contribution tips page, and the project publishes a development roadmap on its documentation site, so the direction of travel is visible even when the source does not spell it out.

## Licence, releases and the missing support commitment

The licence needs a second look before you build on this. The manifest declares license = "Apache-2.0" with the LICENSE file listed under license-files, and the project also carries a NOTICE-style file at the top level of the repository alongside a CONTEXT.md and a CONTRIBUTING.md. GitHub classifies the licence of this repository as a custom licence it cannot categorise, which is a signal to read the LICENSE file and the third party data directory yourself rather than assume the Apache 2.0 grant covers everything you find in the tree.

Third party content matters here. The README states that a single CHEASE equilibrium file for testing and demonstration purposes lives in the data/third_party/geo directory, that it comes from PINT, and that a PINT licence file sits in the same directory. Read that licence file before shipping that equilibrium inside a commercial product.

On versions, the manifest takes its version dynamically from an attribute named TORAX_VERSION in the torax package, so the number you install comes from the published artifact rather than from a literal in pyproject.toml. Recent tags include v1.4.1 on 2026-06-16, v1.4.2 on 2026-06-23 and v1.4.3 on 2026-07-03, and the last push to the repository was on 2026-10-20, which is more recent than the newest release tag in that list. The README also states plainly that TORAX is not an officially supported Google product, so the project carries DeepMind authorship without carrying a support commitment.

## Conclusion

TORAX belongs in the hands of fusion researchers who need gradients through a core transport model, who are already running JAX for machine learned surrogates, and who can pin a scientific software stack and read its third party licence files. It is the wrong choice for a quick plasma picture, for a code base that cannot absorb the JAX and Python 3.12 dependency floor, or for a group that needs a supported product rather than a research tool. Before you commit, check three things: that your equilibrium files load through the CHEASE, FBT or EQDSK readers, that the pinned fusion_surrogates 0.4.7 and imas-python 2.1.0 builds install in your environment, and that the LICENSE file at the repository root and the PINT licence in data/third_party/geo match the terms your project needs.

## FAQ

### What does TORAX simulate, and who is it for?

TORAX is a differentiable one dimensional tokamak core transport simulator written in Python with JAX, aimed at forward modelling, pulse design, trajectory optimisation and controller design. It couples heat and particle transport PDEs with current diffusion and is built for workflows that need gradients through the model.

### How do I install TORAX and run the basic example?

Install Python 3.12 or greater, install tkinter, create a virtual environment, then install the package from PyPI with pip install torax. The README suggests checking the install with run_torax --config=examples/basic_config.py --quit, and running the same configuration without the quit flag for a full simulation with a terminal progress bar.

### Which solver and geometry options does TORAX offer?

The README lists a linear solver with Pereverzev-Corrigan terms and a predictor-corrector method, a Newton-Raphson nonlinear solver, and a nonlinear option that uses the jaxopt library. Geometry is supplied through CHEASE, FBT or EQDSK equilibrium files, and time dependent geometry is supported by providing a time series of those files.

### How does TORAX handle error checking and the JAX compilation cache?

Setting the TORAX_ERRORS_ENABLED environment variable to true enables error checking in internal routines for debugging, and its default is false because it is incompatible with the persistent compilation cache. The JAX_COMPILATION_CACHE_DIR variable lets JAX store compiled programs on the filesystem to cut recompilation time in some cases.

## Sources

- [google-deepmind/torax on GitHub](https://github.com/google-deepmind/torax)
- [Issues](https://github.com/google-deepmind/torax/issues)
- [Project website](https://torax.readthedocs.io)
- [README](https://github.com/google-deepmind/torax/blob/main/README.md)
- [Releases](https://github.com/google-deepmind/torax/releases)

---

Hysen Labs editorial analysis, written from the project's own repository and release notes. Cite the canonical page: https://hysenlabs.com/projects/google-deepmind-torax