Switch the newton_raphson_solve_block to use jax_root_finding.root_newton_raphson, and JIT the entire newton_raphson_solve_block.

The `jax` and `jaxlib` versions are incremented to `0.7` in order to support the new `--xla_backend_extra_options=xla_cpu_flatten_after_fusion` flag, and Python incremented to 3.11 to support jaxlib 0.7.0. This flag treats all functions wrapped with a `jax.jit` within a larger `jax.jit` as independently compilable subfunctions (ie. prevents compiler inlining). It is critical for preventing blowups of compile-times (XLA:CPU compile-times scale roughly quadratically with the number of ops otherwise).

Use of jax.lax.custom_root is temporarily disabled until a workaround is found for it significantly increasing compile times.

This increases simulation speed by around 20%, but at the cost of ~10% increases in compile-time

PiperOrigin-RevId: 782947706

Author: sbodenstein

.github/workflows/linting.yml

@@ -14,7 +14,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ['3.10']
+        python-version: ['3.11']
     steps:
     - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}

.github/workflows/pytest.yml

@@ -52,7 +52,7 @@ jobs:
         # Can't reference env variables in matrix
         num-shards: ${{ fromJson(needs.shards-job.outputs.num-shards) }}
         shard-id: ${{ fromJson(needs.shards-job.outputs.shard-ids) }}
-        python-version: ['3.10']
+        python-version: ['3.11']
         os-version: [ubuntu-latest]
 
     steps:

.readthedocs.yaml

@@ -14,7 +14,7 @@ version: 2
 build:
   os: ubuntu-22.04
   tools:
-    python: "3.10"
+    python: "3.11"
   commands:
     - pip install -r docs/requirements.txt
     - cd docs && make html

README.md

@@ -70,7 +70,7 @@ is outlined in our readthedocs pages.
 
 ### Requirements
 
-Install Python 3.10 or greater.
+Install Python 3.11 or greater.
 
 Make sure that tkinter is installed:
 

docs/installation.rst

@@ -8,7 +8,7 @@ Installation Guide
 Requirements
 ============
 
-Install Python 3.10 or greater.
+Install Python 3.11 or greater.
 
 Make sure that tkinter is installed:
 

pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta"
 name = "torax"
 description = "Differentiable 1D tokamak plasma transport simulator in JAX."
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 license = "Apache-2.0"
 license-files = ["LICENSE"]
 authors = [{name = "TORAX authors", email="torax-dev@google.com"}]
@@ -18,14 +18,14 @@ dependencies = [
     "absl-py>=2.0.0",
     "typing_extensions>=4.2.0",
     "immutabledict>=1.0.0",
-    "jax>=0.4.32",
-    "jaxlib>=0.4.32",
+    "jax>=0.7.0",
+    "jaxlib>=0.7.0",
     "jaxopt>=0.8.2",
     "flax>=0.10.0",
     "fusion_surrogates==0.1.0",
     "matplotlib>=3.3.0",
     "numpy>2",
-    "setuptools;python_version>='3.10'",
+    "setuptools;python_version>='3.11'",
     "chex>=0.1.88",
     "equinox>=0.11.3",
     "PyYAML>=6.0.1",

torax/__init__.py

@@ -35,6 +35,12 @@
 
 # pylint: enable=g-importing-member
 
+
+os.environ['XLA_FLAGS'] = (
+    os.environ.get('XLA_FLAGS', '')
+    + ' --xla_backend_extra_options=xla_cpu_flatten_after_fusion'
+)
+
 __version__ = version.TORAX_VERSION
 __version_info__ = version.TORAX_VERSION_INFO
 

torax/_src/fvm/jax_root_finding.py

@@ -47,6 +47,7 @@ def root_newton_raphson(
     delta_reduction_factor: float = 0.5,
     tau_min: float = 0.01,
     log_iterations: bool = False,
+    use_jax_custom_root: bool = True,
 ) -> tuple[jax.Array, RootMetadata]:
   """A differentiable Newton-Raphson root finder.
 
@@ -66,6 +67,9 @@ def root_newton_raphson(
       routine resets at a lower timestep.
     log_iterations: If true, output diagnostic information from within iteration
       loop.
+    use_jax_custom_root: If true, use jax.lax.custom_root to allow for
+      differentiable solving. This can increase compile times even when no
+      derivatives are requested.
 
   Returns:
     A tuple `(x_root, RootMetadata(...))`.
@@ -113,13 +117,16 @@ def _newton_raphson(f, x):
   def back(g, y):
     return jnp.linalg.solve(jax.jacfwd(g)(y), y)
 
-  x_out, metadata = jax.lax.custom_root(
-      f=fun,
-      initial_guess=x0,
-      solve=_newton_raphson,
-      tangent_solve=back,
-      has_aux=True,
-  )
+  if use_jax_custom_root:
+    x_out, metadata = jax.lax.custom_root(
+        f=fun,
+        initial_guess=x0,
+        solve=_newton_raphson,
+        tangent_solve=back,
+        has_aux=True,
+    )
+  else:
+    x_out, metadata = _newton_raphson(fun, x0)
 
   # Tell the caller whether or not x_new successfully reduces the residual below
   # the tolerance by providing an extra output, error.