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

pyproject.toml

@@ -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.