Add dependency on typeguard.

PiperOrigin-RevId: 793588618

Author: sbodenstein

pyproject.toml

@@ -33,13 +33,14 @@ dependencies = [
     "netcdf4>=1.7.2",
     "h5netcdf>=1.3.0",
     "scipy>=1.13.0",
-    "jaxtyping>=0.2.28",
+    "jaxtyping>=0.3.2",
     "contourpy>=1.2.1",
     "eqdsk>=0.4.0",
     "pydantic>=2.10.5",
     "tqdm>=4.67.0",
     "treelib>=1.3.2",
     "imas-python>=2.0.1",
+    "typeguard>=4.4.4",
 ]
 
 dynamic = ["version"]

torax/_src/array_typing.py

@@ -13,11 +13,19 @@
 # limitations under the License.
 # ============================================================================
 """Common types for using jaxtyping in TORAX."""
-import chex
+from typing import TypeAlias
+import jax
 import jaxtyping as jt
+import numpy as np
+import typeguard
 
-ScalarFloat = jt.Float[chex.Array | float, ""]
-ScalarBool = jt.Bool[chex.Array | bool, ""]
-ScalarInt = jt.Int[chex.Array | int, ""]
-ArrayFloat = jt.Float[chex.Array, "rhon"]
-ArrayBool = jt.Bool[chex.Array, "rhon"]
+Array: TypeAlias = jax.Array | np.ndarray
+
+ScalarFloat = jt.Float[Array | float, ""]
+ScalarBool = jt.Bool[Array | bool, ""]
+ScalarInt = jt.Int[Array | int, ""]
+
+ArrayFloat = jt.Float[Array, "rhon"]
+ArrayBool = jt.Bool[Array, "rhon"]
+
+jaxtyped = jt.jaxtyped(typechecker=typeguard.typechecked)

torax/_src/math_utils.py

@@ -111,7 +111,12 @@ def cell_to_face(
   return face_values
 
 
-def tridiag(diag: jax.Array, above: jax.Array, below: jax.Array) -> jax.Array:
+@array_typing.jaxtyped
+def tridiag(
+    diag: jt.Shaped[array_typing.Array, 'size'],
+    above: jt.Shaped[array_typing.Array, 'size-1'],
+    below: jt.Shaped[array_typing.Array, 'size-1'],
+) -> jt.Shaped[array_typing.Array, 'size size']:
   """Builds a tridiagonal matrix.
 
   Args: