Add helper function for counting number of times a function has been JIT compiled.

Creating a standalone function allows us to:
- have single point of responsibility
- test whether the assumptions we make on JAX internals are valid.

PiperOrigin-RevId: 743897133

Author: Nush395

torax/jax_utils.py

@@ -18,6 +18,7 @@
 import functools
 import os
 from typing import Any, Callable, Optional, TypeVar
+
 import chex
 import equinox as eqx
 import jax
@@ -257,4 +258,29 @@ def py_cond(
     return false_fun()
 
 
+def get_number_of_compiles(
+    jitted_function: Callable[..., Any],
+) -> int:
+  """Helper function for debugging JAX compilation.
+
+  This counts the number of times the function has been JIT compiled. This does
+  not include any uses of the AOT compile workflow.
+
+  Args:
+    jitted_function: A function that has been wrapped with `jax.jit`.
+  Returns:
+    The number of times the function has been compiled.
+  Raises:
+    RuntimeError: If the function does not have a _cache_size attribute.
+  """
+  # pylint: disable=protected-access
+  if not hasattr(jitted_function, '_cache_size'):
+    raise RuntimeError(
+        'The function does not have a _cache_size attribute. Possibly because'
+        ' the function was not jitted.'
+    )
+  return jitted_function._cache_size()
+  # pylint: enable=protected-access
+
+
 # pylint: enable=g-bare-generic

torax/tests/jax_utils_test.py

@@ -16,6 +16,7 @@
 from unittest import mock
 from absl.testing import absltest
 from absl.testing import parameterized
+import jax
 from jax import numpy as jnp
 from torax import jax_utils
 
@@ -105,6 +106,28 @@ def test_f32_int_dtype(self):
     """Test that the dtype is int32 when JAX_PRECISION is set to 'f32'."""
     self.assertEqual(jax_utils.get_int_dtype(), jnp.int32)
 
+  def test_get_number_of_compiles(self):
+    """Check assumptions on JAX internals are valid."""
+
+    def f(x: jax.Array):
+      return x
+
+    jit_f = jax.jit(f)
+    self.assertTrue(hasattr(jit_f, '_cache_size'))
+    # Should be 0 before any calls.
+    self.assertEqual(jax_utils.get_number_of_compiles(jit_f), 0)
+
+    # Should be 1 after one call.
+    jit_f(jnp.array(0))
+    self.assertEqual(jax_utils.get_number_of_compiles(jit_f), 1)
+    # Should be 1 after another call with same shape.
+    jit_f(jnp.array(1))
+    self.assertEqual(jax_utils.get_number_of_compiles(jit_f), 1)
+
+    # Should be 2 after another call with different shape.
+    jit_f(jnp.array([1]))
+    self.assertEqual(jax_utils.get_number_of_compiles(jit_f), 2)
+
 
 if __name__ == '__main__':
   absltest.main()