Use Python 3.12 type annotation.

torax/_src/array_typing.py

@@ -14,7 +14,7 @@
 # ============================================================================
 """Common types for using jaxtyping in TORAX."""
 
-from typing import TypeAlias, TypeVar
+from typing import TypeVar
 import jax
 import jaxtyping as jt
 import numpy as np
@@ -23,21 +23,21 @@
 
 T = TypeVar("T")
 
-Array: TypeAlias = jax.Array | np.ndarray
-
-FloatScalar: TypeAlias = jt.Float[Array | float, ""]
-BoolScalar: TypeAlias = jt.Bool[Array | bool, ""]
-IntScalar: TypeAlias = jt.Int[Array | int, ""]
-
-FloatVector: TypeAlias = jt.Float[Array, "_"]
-BoolVector: TypeAlias = jt.Bool[Array, "_"]
-IntVector: TypeAlias = jt.Int[Array, "_"]
-FloatVectorCell: TypeAlias = jt.Float[Array, "rhon"]
-FloatVectorCellPlusBoundaries: TypeAlias = jt.Float[Array, "rhon+2"]
-FloatMatrixCell: TypeAlias = jt.Float[Array, "rhon rhon"]
-FloatVectorFace: TypeAlias = jt.Float[Array, "rhon+1"]
-BoolVectorCell: TypeAlias = jt.Bool[Array, "rhon"]
-BoolVectorFace: TypeAlias = jt.Bool[Array, "rhon+1"]
+type Array = jax.Array | np.ndarray
+
+type FloatScalar = jt.Float[Array.__value__ | float, ""]
+type BoolScalar = jt.Bool[Array.__value__ | bool, ""]
+type IntScalar = jt.Int[Array.__value__ | int, ""]
+
+type FloatVector = jt.Float[Array.__value__, "_"]
+type BoolVector = jt.Bool[Array.__value__, "_"]
+type IntVector = jt.Int[Array.__value__, "_"]
+type FloatVectorCell = jt.Float[Array.__value__, "rhon"]
+type FloatVectorCellPlusBoundaries = jt.Float[Array.__value__, "rhon+2"]
+type FloatMatrixCell = jt.Float[Array.__value__, "rhon rhon"]
+type FloatVectorFace = jt.Float[Array.__value__, "rhon+1"]
+type BoolVectorCell = jt.Bool[Array.__value__, "rhon"]
+type BoolVectorFace = jt.Bool[Array.__value__, "rhon+1"]
 
 
 def jaxtyped(fn: T) -> T:

torax/_src/config/build_runtime_params.py

@@ -22,7 +22,7 @@
 """
 
 import dataclasses
-from typing import Any, Callable, Mapping, Sequence, TypeAlias
+from typing import Any, Callable, Mapping, Sequence
 
 import chex
 import equinox as eqx
@@ -54,12 +54,12 @@
 # pylint: disable=invalid-name
 
 
-ReplaceablePytreeNodes: TypeAlias = (
+type ReplaceablePytreeNodes = (
     interpolated_param_1d.TimeVaryingScalar
     | interpolated_param_2d.TimeVaryingArray
     | chex.Numeric
 )
-ValidUpdates: TypeAlias = (
+type ValidUpdates = (
     interpolated_param_1d.TimeVaryingScalarUpdate
     | interpolated_param_2d.TimeVaryingArrayUpdate
     | chex.Numeric

torax/_src/config/config_loader.py

@@ -20,20 +20,20 @@
 import sys
 import types
 import typing
-from typing import Any, Literal, TypeAlias
+from typing import Any, Literal
 
 from torax._src import path_utils
 from torax._src.plotting import plotruns_lib
 from torax._src.torax_pydantic import model_config
 
-ExampleConfig: TypeAlias = Literal[
+type ExampleConfig = Literal[
     'basic_config',
     'iterhybrid_predictor_corrector',
     'iterhybrid_rampup',
     'step_flattop_bgb',
 ]
 
-ExamplePlotConfig: TypeAlias = Literal[
+type ExamplePlotConfig = Literal[
     'default_plot_config',
     'global_params_plot_config',
     'simple_plot_config',
@@ -52,7 +52,7 @@ def _get_path(path):
     assert path.is_file(), f'Path {path} to the example config does not exist.'
     return path
 
-  return {path: _get_path(path) for path in typing.get_args(ExampleConfig)}
+  return {path: _get_path(path) for path in typing.get_args(ExampleConfig.__value__)}
 
 
 def example_plot_config_paths() -> dict[ExamplePlotConfig, pathlib.Path]:
@@ -66,7 +66,7 @@ def _get_path(path):
     assert path.is_file(), f'Path {path} to the example config does not exist.'
     return path
 
-  return {path: _get_path(path) for path in typing.get_args(ExamplePlotConfig)}
+  return {path: _get_path(path) for path in typing.get_args(ExamplePlotConfig.__value__)}
 
 
 # Taken from

torax/_src/config/runtime_validation_utils.py

@@ -16,7 +16,7 @@
 
 from collections.abc import Mapping
 import logging
-from typing import Annotated, Any, Final, TypeAlias
+from typing import Annotated, Any, Final
 
 import numpy as np
 import pydantic
@@ -35,7 +35,7 @@ def time_varying_array_defined_at_1(
   return time_varying_array
 
 
-TimeVaryingArrayDefinedAtRightBoundaryAndBounded: TypeAlias = Annotated[
+type TimeVaryingArrayDefinedAtRightBoundaryAndBounded = Annotated[
     torax_pydantic.TimeVaryingArray,
     torax_pydantic.array_bounds_validator(ge=1.0),
     pydantic.AfterValidator(time_varying_array_defined_at_1),
@@ -82,7 +82,7 @@ def _ion_mixture_after_validator(
   return value
 
 
-IonMapping: TypeAlias = Annotated[
+type IonMapping = Annotated[
     Mapping[str, torax_pydantic.TimeVaryingScalar],
     pydantic.BeforeValidator(_ion_mixture_before_validator),
     pydantic.AfterValidator(_ion_mixture_after_validator),

torax/_src/config/tests/config_loader_test.py

@@ -28,7 +28,7 @@ class ConfigLoaderTest(parameterized.TestCase):
   def test_example_config_paths(self):
     self.assertLen(
         config_loader.example_config_paths(),
-        len(typing.get_args(config_loader.ExampleConfig)),
+        len(typing.get_args(config_loader.ExampleConfig.__value__)),
     )
 
   @parameterized.product(

torax/_src/core_profiles/plasma_composition/impurity_fractions.py

@@ -15,7 +15,7 @@
 """Ion mixture model and impurity fractions model for plasma composition."""
 from collections.abc import Mapping
 import dataclasses
-from typing import Annotated, Any, Literal, TypeAlias
+from typing import Annotated, Any, Literal
 
 import chex
 import jax
@@ -103,7 +103,7 @@ def _impurity_after_validator(
   return value
 
 
-ImpurityMapping: TypeAlias = Annotated[
+type ImpurityMapping = Annotated[
     Mapping[str, torax_pydantic.NonNegativeTimeVaryingArray],
     pydantic.BeforeValidator(_impurity_before_validator),
     pydantic.AfterValidator(_impurity_after_validator),

torax/_src/edge/extended_lengyel_standalone.py

@@ -134,11 +134,11 @@ def get_unique_roots(self) -> ExtendedLengyelOutputs | None:
       if field.name in ['roots', 'multiple_roots_found', 'solver_status']:
         continue  # Skip recursive field and internal flags
       value = getattr(roots, field.name)
-      if isinstance(value, array_typing.Array):
+      if isinstance(value, jax.Array):
         fields_to_compress[field.name] = jnp.asarray(value)
       elif isinstance(value, Mapping):
         for k, v in value.items():
-          if isinstance(v, array_typing.Array):
+          if isinstance(v, jax.Array):
             fields_to_compress[f'{field.name}_{k}'] = jnp.asarray(v)
 
     ref_shape = jnp.asarray(roots.T_e_target).shape

torax/_src/fvm/block_1d_coeffs.py

@@ -21,7 +21,7 @@
 """
 
 import dataclasses
-from typing import Any, TypeAlias
+from typing import Any
 
 import jax
 
@@ -32,11 +32,11 @@
 # ((a, b), (c, d)) where a, b, c, d are each jax.Array
 #
 # ((a, None), (None, d)) : represents a diagonal block matrix
-OptionalTupleMatrix: TypeAlias = tuple[tuple[jax.Array | None, ...], ...] | None
+type OptionalTupleMatrix = tuple[tuple[jax.Array | None, ...], ...] | None
 
 
 # Alias for better readability.
-AuxiliaryOutput: TypeAlias = Any
+type AuxiliaryOutput = Any
 
 
 @jax.tree_util.register_dataclass

torax/_src/fvm/discrete_system.py

@@ -24,8 +24,6 @@
 each step is expressed using a matrix multiply.
 """
 
-from typing import TypeAlias
-
 import jax
 from jax import numpy as jnp
 from torax._src import tridiagonal
@@ -34,8 +32,8 @@
 from torax._src.fvm import convection_terms
 from torax._src.fvm import diffusion_terms
 
-AuxiliaryOutput: TypeAlias = block_1d_coeffs.AuxiliaryOutput
-Block1DCoeffs: TypeAlias = block_1d_coeffs.Block1DCoeffs
+type AuxiliaryOutput = block_1d_coeffs.AuxiliaryOutput
+type Block1DCoeffs = block_1d_coeffs.Block1DCoeffs
 
 
 def calc_c(

torax/_src/fvm/optimizer_solve_block.py

@@ -16,8 +16,6 @@
 See function docstring for details.
 """
 
-from typing import TypeAlias
-
 import jax
 import jax.numpy as jnp
 from torax._src import jax_utils
@@ -36,7 +34,7 @@
 from torax._src.solver import predictor_corrector_method
 from torax._src.sources import source_profiles
 
-AuxiliaryOutput: TypeAlias = block_1d_coeffs.AuxiliaryOutput
+type AuxiliaryOutput = block_1d_coeffs.AuxiliaryOutput
 
 
 @jax.jit(

torax/_src/fvm/residual_and_loss.py

@@ -21,7 +21,6 @@
 """
 
 import functools
-from typing import TypeAlias
 
 import chex
 import jax
@@ -42,7 +41,7 @@
 from torax._src.pedestal_model import pedestal_transition_state as pedestal_transition_state_lib
 from torax._src.sources import source_profiles
 
-Block1DCoeffs: TypeAlias = block_1d_coeffs.Block1DCoeffs
+type Block1DCoeffs = block_1d_coeffs.Block1DCoeffs
 
 
 @jax.jit(

torax/_src/geometry/fbt.py

@@ -18,7 +18,7 @@
 import logging
 from typing import Annotated
 from typing import Any
-from typing import Literal, TypeAlias
+from typing import Literal
 
 import jax
 import numpy as np
@@ -34,7 +34,7 @@
 import typing_extensions
 
 # pylint: disable=invalid-name
-LY_OBJECT_TYPE: TypeAlias = (
+type LY_OBJECT_TYPE = (
     str | Mapping[str, torax_pydantic.NumpyArray | float]
 )
 

torax/_src/geometry/geometry.py

@@ -423,7 +423,7 @@ def stack_geometries(geometries: Sequence[GeometryT]) -> GeometryT:
     field_name = field.name
     field_value = getattr(first_geo, field_name)
     # Stack stackable fields. Save first geo's value for non-stackable fields.
-    if isinstance(field_value, (array_typing.Array, array_typing.FloatScalar)):  # pyrefly: ignore[invalid-argument]
+    if isinstance(field_value, (jax.Array, np.ndarray, float)):
       field_values = [getattr(geo, field_name) for geo in geometries]
       stacked_data[field_name] = np.stack(field_values)
     else:

torax/_src/geometry/tests/eqdsk_test.py

@@ -47,7 +47,7 @@ def test_eqdsk_cocos_conversion_is_consistent(self):
       name = field.name
       val1 = getattr(geo_cocos2, name)
       val2 = getattr(geo_cocos11, name)
-      if isinstance(val1, array_typing.Array):
+      if isinstance(val1, np.ndarray):
         np.testing.assert_allclose(
             val1, val2, err_msg=f'Field "{name}" mismatch.'
         )
@@ -75,7 +75,7 @@ def test_build_geometry_from_eqdsk_object(self):
       name = field.name
       val1 = getattr(geo_file, name)
       val2 = getattr(geo_obj, name)
-      if isinstance(val1, array_typing.Array):
+      if isinstance(val1, np.ndarray):
         np.testing.assert_allclose(
             val1, val2, err_msg=f'Field "{name}" mismatch.'
         )
@@ -107,7 +107,7 @@ def test_eqdsk_serialization_round_trip(self):
       name = field.name
       val1 = getattr(geo_original, name)
       val2 = getattr(geo_restored, name)
-      if isinstance(val1, array_typing.Array):
+      if isinstance(val1, np.ndarray):
         np.testing.assert_allclose(
             val1, val2, err_msg=f'Field "{name}" mismatch (dict).'
         )

torax/_src/interpolated_param.py

@@ -17,7 +17,7 @@
 import abc
 from collections.abc import Mapping
 import enum
-from typing import Final, Literal, TypeAlias
+from typing import Final, Literal
 import chex
 import jax
 import jax.numpy as jnp
@@ -69,23 +69,23 @@ class InterpolationMode(enum.Enum):
   STEP = 'step'
 
 
-InterpolationModeLiteral: TypeAlias = Literal[
+type InterpolationModeLiteral = Literal[
     'step', 'STEP', 'piecewise_linear', 'PIECEWISE_LINEAR'
 ]
 
 
-_ArrayOrListOfFloats: TypeAlias = array_typing.Array | list[float]
+type _ArrayOrListOfFloats = array_typing.Array | list[float]
 
 # Config input types convertible to InterpolatedParam objects.
-InterpolatedVarSingleAxisInput: TypeAlias = (
+type InterpolatedVarSingleAxisInput = (
     float
     | dict[float, float]
     | bool
     | dict[float, bool]
     | tuple[_ArrayOrListOfFloats, _ArrayOrListOfFloats]
     | xr.DataArray
 )
-InterpolatedVarTimeRhoInput: TypeAlias = (
+type InterpolatedVarTimeRhoInput = (
     # Mapping from time to rho, value interpolated in rho
     Mapping[float, InterpolatedVarSingleAxisInput]
     | float
@@ -98,12 +98,12 @@ class InterpolationMode(enum.Enum):
 # Type-alias for a variable (in rho_norm) to be interpolated in time.
 # If a string is provided, it is assumed to be an InterpolationMode else, the
 # default piecewise linear interpolation is used.
-TimeInterpolatedInput: TypeAlias = (
+type TimeInterpolatedInput = (
     InterpolatedVarSingleAxisInput
     | tuple[InterpolatedVarSingleAxisInput, InterpolationModeLiteral]
 )
 # Type-alias for a variable to be interpolated in time and rho_norm.
-TimeRhoInterpolatedInput: TypeAlias = (
+type TimeRhoInterpolatedInput = (
     InterpolatedVarTimeRhoInput
     | tuple[
         InterpolatedVarTimeRhoInput,

torax/_src/jax_utils.py

@@ -17,7 +17,7 @@
 import contextlib
 import functools
 import os
-from typing import Any, Callable, Literal, ParamSpec, TypeAlias, TypeVar
+from typing import Any, Callable, Literal, ParamSpec, TypeVar
 import chex
 import jax
 from jax import numpy as jnp
@@ -27,9 +27,9 @@
 from packaging import version
 
 T = TypeVar('T')
-BooleanNumeric: TypeAlias = Any  # A bool, or a Boolean array.
+type BooleanNumeric = Any  # A bool, or a Boolean array.
 _State = ParamSpec('_State')
-PyTree: TypeAlias = Any
+type PyTree = Any
 
 _WHILE_LOOP_COUNT_DTYPE = jnp.int32