diff --git a/torax/_src/array_typing.py b/torax/_src/array_typing.py
index f79d0f144..e137d80d9 100644
--- a/torax/_src/array_typing.py
+++ b/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:
diff --git a/torax/_src/config/build_runtime_params.py b/torax/_src/config/build_runtime_params.py
index 5fc382cc3..4761d1b22 100644
--- a/torax/_src/config/build_runtime_params.py
+++ b/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
diff --git a/torax/_src/config/config_loader.py b/torax/_src/config/config_loader.py
index 892918eb2..23ffa643b 100644
--- a/torax/_src/config/config_loader.py
+++ b/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
diff --git a/torax/_src/config/runtime_validation_utils.py b/torax/_src/config/runtime_validation_utils.py
index a68737899..9f73a79be 100644
--- a/torax/_src/config/runtime_validation_utils.py
+++ b/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),
diff --git a/torax/_src/config/tests/config_loader_test.py b/torax/_src/config/tests/config_loader_test.py
index b2ece71e7..0bd03345b 100644
--- a/torax/_src/config/tests/config_loader_test.py
+++ b/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(
diff --git a/torax/_src/core_profiles/plasma_composition/impurity_fractions.py b/torax/_src/core_profiles/plasma_composition/impurity_fractions.py
index 12692b6a9..3757d18a4 100644
--- a/torax/_src/core_profiles/plasma_composition/impurity_fractions.py
+++ b/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),
diff --git a/torax/_src/edge/extended_lengyel_standalone.py b/torax/_src/edge/extended_lengyel_standalone.py
index ee7f52e4f..4b013e097 100644
--- a/torax/_src/edge/extended_lengyel_standalone.py
+++ b/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
diff --git a/torax/_src/fvm/block_1d_coeffs.py b/torax/_src/fvm/block_1d_coeffs.py
index d828cd5e1..be6d9f5c3 100644
--- a/torax/_src/fvm/block_1d_coeffs.py
+++ b/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
diff --git a/torax/_src/fvm/discrete_system.py b/torax/_src/fvm/discrete_system.py
index e28b7683c..77fe930c8 100644
--- a/torax/_src/fvm/discrete_system.py
+++ b/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(
diff --git a/torax/_src/fvm/optimizer_solve_block.py b/torax/_src/fvm/optimizer_solve_block.py
index bc3388d63..a33f1215a 100644
--- a/torax/_src/fvm/optimizer_solve_block.py
+++ b/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(
diff --git a/torax/_src/fvm/residual_and_loss.py b/torax/_src/fvm/residual_and_loss.py
index 7aa7812b8..9b3bba697 100644
--- a/torax/_src/fvm/residual_and_loss.py
+++ b/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(
diff --git a/torax/_src/geometry/fbt.py b/torax/_src/geometry/fbt.py
index 5b838b0ad..f1f0e3fe2 100644
--- a/torax/_src/geometry/fbt.py
+++ b/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]
 )
 
diff --git a/torax/_src/geometry/geometry.py b/torax/_src/geometry/geometry.py
index 8307e4360..dca9003d7 100644
--- a/torax/_src/geometry/geometry.py
+++ b/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:
diff --git a/torax/_src/geometry/tests/eqdsk_test.py b/torax/_src/geometry/tests/eqdsk_test.py
index e12d302f1..b8041ced2 100644
--- a/torax/_src/geometry/tests/eqdsk_test.py
+++ b/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).'
         )
diff --git a/torax/_src/interpolated_param.py b/torax/_src/interpolated_param.py
index fc33f2647..91914f9ed 100644
--- a/torax/_src/interpolated_param.py
+++ b/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,15 +69,15 @@ 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
@@ -85,7 +85,7 @@ class InterpolationMode(enum.Enum):
     | 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,
diff --git a/torax/_src/jax_utils.py b/torax/_src/jax_utils.py
index f64910d3f..deefa20fb 100644
--- a/torax/_src/jax_utils.py
+++ b/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
 
diff --git a/torax/_src/orchestration/jit_run_loop.py b/torax/_src/orchestration/jit_run_loop.py
index 7ce25c0c3..9e8895da2 100644
--- a/torax/_src/orchestration/jit_run_loop.py
+++ b/torax/_src/orchestration/jit_run_loop.py
@@ -14,7 +14,7 @@
 
 """JITted run_loop for iterating over the simulation step function."""
 
-from typing import Any, TypeAlias
+from typing import Any
 import chex
 import jax
 import jax.numpy as jnp
@@ -27,7 +27,7 @@
 from torax._src.orchestration import step_function
 from torax._src.output_tools import post_processing
 
-PyTree: TypeAlias = Any
+type PyTree = Any
 
 
 @jax.jit(static_argnames='max_steps')
diff --git a/torax/_src/output_tools/output.py b/torax/_src/output_tools/output.py
index e369de86b..b8e704228 100644
--- a/torax/_src/output_tools/output.py
+++ b/torax/_src/output_tools/output.py
@@ -776,7 +776,7 @@ def _save_geometry(
               and field_name.removesuffix("_face") in geometry_attributes
           )
           or field_name in _EXCLUDED_GEOMETRY_FIELDS
-          or not isinstance(data, array_typing.Array)
+          or not isinstance(data, (jax.Array, np.ndarray))
       ):
         continue
       if f"{field_name}_face" in geometry_attributes:
diff --git a/torax/_src/pedestal_model/pydantic_model.py b/torax/_src/pedestal_model/pydantic_model.py
index 755ae1819..a93e58ea2 100644
--- a/torax/_src/pedestal_model/pydantic_model.py
+++ b/torax/_src/pedestal_model/pydantic_model.py
@@ -16,7 +16,7 @@
 
 import abc
 import copy
-from typing import Annotated, Any, Literal, TypeAlias
+from typing import Annotated, Any, Literal
 import chex
 import pydantic
 from torax._src import array_typing
@@ -209,9 +209,9 @@ def build_runtime_params(
     )
 
 
-# For new formation and saturation models, add to these TypeAliases via Union.
-FormationConfig: TypeAlias = DelabieScalingFormation | MartinScalingFormation
-SaturationConfig: TypeAlias = ProfileValueSaturation
+# For new formation and saturation models, add to these type aliases via Union.
+type FormationConfig = DelabieScalingFormation | MartinScalingFormation
+type SaturationConfig = ProfileValueSaturation
 
 
 class BasePedestal(torax_pydantic.BaseModelFrozen, abc.ABC):
diff --git a/torax/_src/solver/jax_fixed_point.py b/torax/_src/solver/jax_fixed_point.py
index 17373ddae..211df9800 100644
--- a/torax/_src/solver/jax_fixed_point.py
+++ b/torax/_src/solver/jax_fixed_point.py
@@ -14,13 +14,13 @@
 
 """JAX fixed point functions."""
 
-from typing import Any, Callable, TypeAlias
+from typing import Any, Callable
 import jax
 import jax.numpy as jnp
 from torax._src import jax_utils
 from torax._src.solver import linesearch
 
-PyTree: TypeAlias = Any
+type PyTree = Any
 
 
 def fixed_point(
diff --git a/torax/_src/torax_pydantic/interpolated_param_1d.py b/torax/_src/torax_pydantic/interpolated_param_1d.py
index fee13a0b4..fe30d8270 100644
--- a/torax/_src/torax_pydantic/interpolated_param_1d.py
+++ b/torax/_src/torax_pydantic/interpolated_param_1d.py
@@ -16,7 +16,7 @@
 
 import dataclasses
 import functools
-from typing import Any, TypeAlias
+from typing import Any
 
 import chex
 import equinox as eqx
@@ -280,15 +280,15 @@ def scalar_bounds_validator(
   )
 
 
-PositiveTimeVaryingScalar: TypeAlias = typing_extensions.Annotated[
+type PositiveTimeVaryingScalar = typing_extensions.Annotated[
     TimeVaryingScalar, scalar_bounds_validator(gt=0.0)
 ]
-NonNegativeTimeVaryingScalar: TypeAlias = typing_extensions.Annotated[
+type NonNegativeTimeVaryingScalar = typing_extensions.Annotated[
     TimeVaryingScalar, scalar_bounds_validator(ge=0.0)
 ]
-NonNegativeTimeVaryingScalarStep: TypeAlias = typing_extensions.Annotated[
+type NonNegativeTimeVaryingScalarStep = typing_extensions.Annotated[
     TimeVaryingScalarStep, scalar_bounds_validator(ge=0.0)
 ]
-UnitIntervalTimeVaryingScalar: TypeAlias = typing_extensions.Annotated[
+type UnitIntervalTimeVaryingScalar = typing_extensions.Annotated[
     TimeVaryingScalar, scalar_bounds_validator(ge=0.0, le=1.0)
 ]
diff --git a/torax/_src/torax_pydantic/interpolated_param_2d.py b/torax/_src/torax_pydantic/interpolated_param_2d.py
index 02f42a525..ee1ee4820 100644
--- a/torax/_src/torax_pydantic/interpolated_param_2d.py
+++ b/torax/_src/torax_pydantic/interpolated_param_2d.py
@@ -17,7 +17,7 @@
 from collections.abc import Mapping
 import dataclasses
 import functools
-from typing import Any, Literal, TypeAlias
+from typing import Any, Literal
 
 import chex
 import equinox as eqx
@@ -34,7 +34,7 @@
 import typing_extensions
 import xarray as xr
 
-ValueType: TypeAlias = dict[
+type ValueType = dict[
     float,
     tuple[pydantic_types.NumpyArray1DUnitInterval, pydantic_types.NumpyArray1D],
 ]
@@ -497,7 +497,7 @@ def _conform_data(
     elif isinstance(data, tuple):
       values = []
       for v in data:
-        if isinstance(v, array_typing.Array):
+        if isinstance(v, (jax.Array, np.ndarray)):
           values.append(v)
         elif isinstance(v, list):
           values.append(np.asarray(v))
@@ -755,7 +755,7 @@ def array_bounds_validator(
   )
 
 
-PositiveTimeVaryingArray: TypeAlias = typing_extensions.Annotated[
+type PositiveTimeVaryingArray = typing_extensions.Annotated[
     TimeVaryingArray, array_bounds_validator(gt=0.0)
 ]
 
@@ -905,6 +905,6 @@ def get_face_centers(nx: int, dx: float | None = None) -> np.ndarray:
   return np.linspace(0, nx * dx, nx + 1)
 
 
-NonNegativeTimeVaryingArray: TypeAlias = typing_extensions.Annotated[
+type NonNegativeTimeVaryingArray = typing_extensions.Annotated[
     TimeVaryingArray, array_bounds_validator(ge=0.0)
 ]
diff --git a/torax/_src/torax_pydantic/model_base.py b/torax/_src/torax_pydantic/model_base.py
index 55521dd6e..ad56d6b9a 100644
--- a/torax/_src/torax_pydantic/model_base.py
+++ b/torax/_src/torax_pydantic/model_base.py
@@ -17,7 +17,7 @@
 from collections.abc import Set
 import functools
 import inspect
-from typing import Any, Final, Mapping, Sequence, TypeAlias
+from typing import Any, Final, Mapping, Sequence
 
 import jax
 import pydantic
@@ -27,8 +27,8 @@
 TIME_INVARIANT: Final[str] = '_pydantic_time_invariant_field'
 JAX_STATIC: Final[str] = '_pydantic_jax_static_field'
 
-StaticKwargs: TypeAlias = dict[str, Any]
-DynamicArgs: TypeAlias = list[Any]
+type StaticKwargs = dict[str, Any]
+type DynamicArgs = list[Any]
 
 
 class BaseModelFrozen(pydantic.BaseModel):
diff --git a/torax/_src/torax_pydantic/pydantic_types.py b/torax/_src/torax_pydantic/pydantic_types.py
index a63a10604..f68177570 100644
--- a/torax/_src/torax_pydantic/pydantic_types.py
+++ b/torax/_src/torax_pydantic/pydantic_types.py
@@ -14,22 +14,22 @@
 
 """Pydantic custom types."""
 
-from typing import Annotated, TypeAlias
+from typing import Annotated
 
 import numpy as np
 import pydantic
 
-DataTypes: TypeAlias = float | int | bool
-DtypeName: TypeAlias = str
+type DataTypes = float | int | bool
+type DtypeName = str
 
-NestedList: TypeAlias = (
+type NestedList = (
     DataTypes
     | list[DataTypes]
     | list[list[DataTypes]]
     | list[list[list[DataTypes]]]
 )
 
-NumpySerialized: TypeAlias = tuple[DtypeName, NestedList]
+type NumpySerialized = tuple[DtypeName, NestedList]
 
 
 def _numpy_array_before_validator(
diff --git a/torax/_src/torax_pydantic/torax_pydantic.py b/torax/_src/torax_pydantic/torax_pydantic.py
index 0d9b7fc6f..26ec8bdb5 100644
--- a/torax/_src/torax_pydantic/torax_pydantic.py
+++ b/torax/_src/torax_pydantic/torax_pydantic.py
@@ -15,7 +15,6 @@
 """Pydantic utilities and base classes."""
 
 import functools
-from typing import TypeAlias
 
 import pydantic
 from torax._src.torax_pydantic import interpolated_param_1d
@@ -29,20 +28,20 @@
 
 # Physical units.
 # keep-sorted start
-CubicMeter: TypeAlias = pydantic.PositiveFloat
-GreenwaldFraction: TypeAlias = pydantic.PositiveFloat
-KiloElectronVolt: TypeAlias = pydantic.PositiveFloat
-Meter: TypeAlias = pydantic.PositiveFloat
-MeterPerSecond: TypeAlias = float
-MeterSquaredPerSecond: TypeAlias = float
-Pascal: TypeAlias = float
-Second: TypeAlias = float
-Tesla: TypeAlias = float
+type CubicMeter = pydantic.PositiveFloat
+type GreenwaldFraction = pydantic.PositiveFloat
+type KiloElectronVolt = pydantic.PositiveFloat
+type Meter = pydantic.PositiveFloat
+type MeterPerSecond = float
+type MeterSquaredPerSecond = float
+type Pascal = float
+type Second = float
+type Tesla = float
 # keep-sorted end
-Density: TypeAlias = CubicMeter | GreenwaldFraction
+type Density = CubicMeter | GreenwaldFraction
 
-UnitInterval: TypeAlias = Annotated[float, pydantic.Field(ge=0.0, le=1.0)]
-OpenUnitInterval: TypeAlias = Annotated[float, pydantic.Field(gt=0.0, lt=1.0)]
+type UnitInterval = Annotated[float, pydantic.Field(ge=0.0, le=1.0)]
+type OpenUnitInterval = Annotated[float, pydantic.Field(gt=0.0, lt=1.0)]
 
 NumpyArray = pydantic_types.NumpyArray
 NumpyArray1D = pydantic_types.NumpyArray1D
diff --git a/torax/_src/transport_model/base_qlknn_model.py b/torax/_src/transport_model/base_qlknn_model.py
index e401e841e..814c7b7da 100644
--- a/torax/_src/transport_model/base_qlknn_model.py
+++ b/torax/_src/transport_model/base_qlknn_model.py
@@ -14,13 +14,12 @@
 """Base class for QLKNN Models."""
 
 import abc
-from typing import TypeAlias
 
 import jax
 from torax._src.transport_model import qualikiz_based_transport_model
 
-ModelOutput: TypeAlias = dict[str, jax.Array]
-InputsAndRanges: TypeAlias = dict[str, dict[str, float]]
+type ModelOutput = dict[str, jax.Array]
+type InputsAndRanges = dict[str, dict[str, float]]
 
 
 class BaseQLKNNModel(abc.ABC):
diff --git a/torax/_src/transport_model/tglf/tglf_transport_model.py b/torax/_src/transport_model/tglf/tglf_transport_model.py
index a34fa16bb..655cec012 100644
--- a/torax/_src/transport_model/tglf/tglf_transport_model.py
+++ b/torax/_src/transport_model/tglf/tglf_transport_model.py
@@ -16,7 +16,7 @@
 
 from concurrent import futures
 import dataclasses
-from typing import Annotated, Any, Literal, TypeAlias
+from typing import Annotated, Any, Literal
 
 from absl import logging
 import chex
@@ -76,7 +76,7 @@
     'n_basis_max': 'NBASIS_MAX',
 }
 
-TGLFSettingsValueTypes: TypeAlias = str | float | int | bool | None
+type TGLFSettingsValueTypes = str | float | int | bool | None
 
 
 @jax.tree_util.register_dataclass