Repository: chromatix-team/chromatix (language: python)

Issue:
Add SAS, vectorial thick samples, and docs from Hackathon 2024
We're merging in a lot of updates and fixes from the 2024 Chromatix Hackathon. This adds SAS propagation, updated documentation and installation instructions, bandlimited angular spectrum (BLAS) with shifted destination sampling, vectorial thick samples, and new examples for some of these features.


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/pyright.yaml b/.github/workflows/pyright.yaml
new file mode 100644
index 0000000..623465a
--- /dev/null
+++ b/.github/workflows/pyright.yaml
@@ -0,0 +1,20 @@
+name: pyright 
+on: [push, pull_request]
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+          cache: 'pip'
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install .[dev]
+      - name: Type check with pyright 
+        run: pyright 
diff --git a/pyproject.toml b/pyproject.toml
index 75f806f..9831155 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -9,11 +9,11 @@ description = "Differentiable wave optics library using JAX"
 readme = "README.md"
 requires-python = ">=3.10"
 license = {text = "MIT"}
-dependencies = ["jax >= 0.4.1", "einops >= 0.6.0", "flax >= 0.6.3", "chex>=0.1.5", "optax >=0.1.4", "scipy >= 1.10.0"]
+dependencies = ["jax >= 0.4.6", "einops >= 0.6.0", "flax >= 0.6.3", "chex>=0.1.86", "optax >=0.1.4", "scipy >= 1.10.0", "typing-extensions >=4.12"]
 version = "0.3.0"
 
 [project.optional-dependencies]
-dev = ["mypy>= 0.991", "pytest>=7.2.0", "ruff >= 0.4.8"]
+dev = ["pyright>=1.1.366", "pytest>=7.2.0", "ruff >= 0.4.8"]
 docs = ["mkdocs >= 1.4.2", "mkdocs-material >= 9.0.6", "mkdocstrings-python >= 0.8.3", "mkdocs-jupyter"]
 
 [tool.ruff]
@@ -24,4 +24,9 @@ ignore = ["E402"]
 extend-select = ["I"]
 
 [tool.ruff.lint.per-file-ignores]
-"__init__.py" = ["F401", "F403"]
\ No newline at end of file
+"__init__.py" = ["F401", "F403"]
+
+
+[tool.pyright]
+include = ["src/"]
+ignore = ["src/chromatix/data"]
\ No newline at end of file
diff --git a/src/chromatix/utils/data.py b/src/chromatix/data/data.py
similarity index 96%
rename from src/chromatix/utils/data.py
rename to src/chromatix/data/data.py
index df7d291..c40fe9e 100644
--- a/src/chromatix/utils/data.py
+++ b/src/chromatix/data/data.py
@@ -1,5 +1,3 @@
-from typing import Optional, Tuple, Union
-
 import numpy as np
 
 try:
@@ -9,10 +7,12 @@
 except ModuleNotFoundError:
     USE_CV2 = False
 
+from chromatix.typing import ArrayLike
+
 
 def sqr_dist_to_line(
-    z: np.ndarray, y: np.ndarray, x: np.ndarray, start: np.ndarray, n: np.ndarray
-) -> np.ndarray:
+    z: ArrayLike, y: ArrayLike, x: ArrayLike, start: ArrayLike, n: ArrayLike
+) -> ArrayLike:
     """
     Returns an array with each pixel being assigned to the square distance to
     that line and an array with the distance along the line.
@@ -66,9 +66,9 @@ def draw_line(
 
 
 def filaments_3d(
-    sz: Tuple[int, int, int],
+    sz: tuple[int, int, int],
     intensity: float = 1.0,
-    radius: Union[float, Tuple[float, float, float]] = 0.8,
+    radius: float | tuple[float, float, float] = 0.8,
     rand_offset: float = 0.05,
     rel_theta: float = 1.0,
     num_filaments: int = 50,
@@ -142,7 +142,7 @@ def filaments_3d(
 
 
 def pollen_3d(
-    sz: Tuple[int, int, int],
+    sz: tuple[int, int, int],
     intensity: float = 1.0,
     radius: float = 0.8,
     dphi: float = 0.0,
@@ -213,7 +213,7 @@ def pollen_3d(
 
 
 def siemens_star(
-    num_pixels: int = 512, num_spokes: int = 32, radius: Optional[int] = None
+    num_pixels: int = 512, num_spokes: int = 32, radius: int | None = None
 ) -> np.ndarray:
     """
     Generates a 2D Siemens star image of shape ``num_pixels``. A single input
@@ -253,7 +253,7 @@ def siemens_star(
 if USE_CV2:
 
     def draw_disks(
-        shape: Tuple[int, int], coordinates: np.ndarray, radius: int, color: int = 255
+        shape: tuple[int, int], coordinates: np.ndarray, radius: int, color: int = 255
     ) -> np.ndarray:
         """
         Create a grayscale image with disks drawn at each provided coordinate.
@@ -277,7 +277,7 @@ def draw_disks(
 else:
 
     def draw_disks(
-        shape: Tuple[int, int], coordinates: np.ndarray, radius: int, color: int = 255
+        shape: tuple[int, int], coordinates: np.ndarray, radius: int, color: int = 255
     ) -> np.ndarray:
         """
         Create a grayscale image with disks drawn at each provided coordinate.
@@ -311,8 +311,8 @@ def __init__(
         N: int,
         num_points: int,
         radius: int,
-        shape: Tuple[int, int],
-        z_range: Tuple[int, int],
+        shape: tuple[int, int],
+        z_range: tuple[int, int],
     ):
         """
         Create a dataset of random 3D coordinates and their associated image.
@@ -419,7 +419,7 @@ def __getitem__(self, idx: int) -> np.ndarray:
             ]  # TODO add weight
             return coords, image
 
-    def __call__(self) -> Tuple[np.ndarray, np.ndarray]:
+    def __call__(self) -> tuple[np.ndarray, np.ndarray]:
         """
         Get a new sample. Automatically iterates through samples of coordinates
         with every call. Will cause the random coordinates to be regenerated
diff --git a/src/chromatix/elements/amplitude_masks.py b/src/chromatix/elements/amplitude_masks.py
index ff4f6dd..e67f48e 100644
--- a/src/chromatix/elements/amplitude_masks.py
+++ b/src/chromatix/elements/amplitude_masks.py
@@ -1,13 +1,14 @@
-from typing import Callable, Tuple, Union
+from typing import Callable
 
-from chex import Array, PRNGKey
+from chex import PRNGKey
 from flax import linen as nn
+from jax import Array
 
 from chromatix.elements.utils import register
-
-from ..field import Field
-from ..functional.amplitude_masks import amplitude_change
-from ..ops import binarize
+from chromatix.field import Field
+from chromatix.functional.amplitude_masks import amplitude_change
+from chromatix.ops import binarize
+from chromatix.typing import ArrayLike
 
 __all__ = ["AmplitudeMask"]
 
@@ -27,7 +28,7 @@ class AmplitudeMask(nn.Module):
         is_binary: binarize the amplitude mask if True.
     """
 
-    amplitude: Union[Array, Callable[[PRNGKey, Tuple[int, int]], Array]]
+    amplitude: ArrayLike | Callable[[PRNGKey, tuple[int, int]], Array]
     is_binary: bool
 
     @nn.compact
diff --git a/src/chromatix/elements/convenience.py b/src/chromatix/elements/convenience.py
index 1210901..7d4728a 100644
--- a/src/chromatix/elements/convenience.py
+++ b/src/chromatix/elements/convenience.py
@@ -1,8 +1,7 @@
-from typing import Optional, Tuple, Union
-
 import flax.linen as nn
 import jax.numpy as jnp
-from chex import Array
+
+from chromatix.typing import ScalarLike
 
 from ..field import Field
 from ..ops import binarize, quantize
@@ -31,8 +30,8 @@ class ScaleAndBias(nn.Module):
     broadcastable to the shape of the incoming ``Field``.
     """
 
-    bias: Union[float, Array]
-    scale: Union[float, Array]
+    bias: ScalarLike
+    scale: ScalarLike
 
     @nn.compact
     def __call__(self, field: Field) -> Field:
@@ -61,7 +60,7 @@ class Quantize(nn.Module):
     """
 
     bit_depth: int
-    range: Optional[Tuple[int, int]] = None
+    range: tuple[int, int] | None = None
 
     @nn.compact
     def __call__(self, field: Field) -> Field:
diff --git a/src/chromatix/elements/lenses.py b/src/chromatix/elements/lenses.py
index fb98b54..c47122f 100644
--- a/src/chromatix/elements/lenses.py
+++ b/src/chromatix/elements/lenses.py
@@ -1,9 +1,11 @@
-from typing import Callable, Optional, Union
+from typing import Callable
 
 import flax.linen as nn
 from chex import PRNGKey
+from jax import Array
 
 from chromatix.elements.utils import register
+from chromatix.typing import ScalarLike
 
 from .. import functional as cf
 from ..field import Field
@@ -29,9 +31,9 @@ class ThinLens(nn.Module):
             to the incoming ``Field``.
     """
 
-    f: Union[float, Callable[[PRNGKey], float]]
-    n: Union[float, Callable[[PRNGKey], float]]
-    NA: Optional[Union[float, Callable[[PRNGKey], float]]] = None
+    f: ScalarLike | Callable[[PRNGKey], Array]
+    n: ScalarLike | Callable[[PRNGKey], Array]
+    NA: ScalarLike | Callable[[PRNGKey], Array] | None = None
 
     @nn.compact
     def __call__(self, field: Field) -> Field:
@@ -60,9 +62,9 @@ class FFLens(nn.Module):
         inverse: Whether to use IFFT (default is False, which uses FFT).
     """
 
-    f: Union[float, Callable[[PRNGKey], float]]
-    n: Union[float, Callable[[PRNGKey], float]]
-    NA: Optional[Union[float, Callable[[PRNGKey], float]]] = None
+    f: ScalarLike | Callable[[PRNGKey], Array]
+    n: ScalarLike | Callable[[PRNGKey], Array]
+    NA: ScalarLike | Callable[[PRNGKey], Array] | None = None
     inverse: bool = False
 
     @nn.compact
@@ -93,10 +95,10 @@ class DFLens(nn.Module):
         inverse: Whether to use IFFT (default is False, which uses FFT).
     """
 
-    d: Union[float, Callable[[PRNGKey], float]]
-    f: Union[float, Callable[[PRNGKey], float]]
-    n: Union[float, Callable[[PRNGKey], float]]
-    NA: Optional[Union[float, Callable[[PRNGKey], float]]] = None
+    d: ScalarLike | Callable[[PRNGKey], Array]
+    f: ScalarLike | Callable[[PRNGKey], Array]
+    n: ScalarLike | Callable[[PRNGKey], Array]
+    NA: ScalarLike | Callable[[PRNGKey], Array] | None = None
     inverse: bool = False
 
     @nn.compact
diff --git a/src/chromatix/elements/phase_masks.py b/src/chromatix/elements/phase_masks.py
index fe72b7d..bb56f28 100644
--- a/src/chromatix/elements/phase_masks.py
+++ b/src/chromatix/elements/phase_masks.py
@@ -1,19 +1,23 @@
-from typing import Callable, Optional, Tuple, Union
+from typing import Callable, Sequence
+
 import jax.numpy as jnp
-from chex import Array, PRNGKey
+from chex import PRNGKey
 from flax import linen as nn
+from jax import Array
 from jax.scipy.ndimage import map_coordinates
+
+from chromatix.elements.utils import register
 from chromatix.field import Field
-from chromatix.functional import wrap_phase, phase_change
-from chromatix.utils import seidel_aberrations, zernike_aberrations
+from chromatix.functional import phase_change, wrap_phase
 from chromatix.ops import quantize
-from chromatix.elements.utils import register
+from chromatix.typing import ArrayLike, ScalarLike
+from chromatix.utils import seidel_aberrations, zernike_aberrations
 
 __all__ = [
-    "PhaseMask",
-    "SpatialLightModulator",
     "SeidelAberrations",
     "ZernikeAberrations",
+    "PhaseMask",
+    "SpatialLightModulator",
 ]
 
 
@@ -47,10 +51,10 @@ class PhaseMask(nn.Module):
         NA: The numerical aperture of the system's objective. Defaults to None.
     """
 
-    phase: Union[Array, Callable[[PRNGKey, Tuple[int, int], float, float], Array]]
-    f: Optional[float] = None
-    n: Optional[float] = None
-    NA: Optional[float] = None
+    phase: ArrayLike | Callable[[PRNGKey, tuple[int, int], Array, Array], Array]
+    f: ScalarLike | None = None
+    n: ScalarLike | None = None
+    NA: ScalarLike | None = None
 
     @nn.compact
     def __call__(self, field: Field) -> Field:
@@ -118,16 +122,15 @@ class SpatialLightModulator(nn.Module):
         n: Refractive index of the system's objective. Defaults to None.
         NA: The numerical aperture of the system's objective. Defaults to None.
     """
-
-    phase: Union[Array, Callable[[PRNGKey, Tuple[int, int], float, float], Array]]
-    shape: Tuple[int, int]
-    spacing: float
-    phase_range: Tuple[float, float]
-    num_bits: Optional[Union[int, float]] = None
+    phase: ArrayLike | Callable[[PRNGKey, tuple[int, int], Array, Array], Array]
+    shape: tuple[int, int]
+    spacing: ScalarLike
+    phase_range: ArrayLike
+    num_bits: int | None = None
     interpolation_order: int = 0
-    f: Optional[float] = None
-    n: Optional[float] = None
-    NA: Optional[float] = None
+    f: ScalarLike | None = None
+    n: ScalarLike | None = None
+    NA: ScalarLike | None = None
 
     @nn.compact
     def __call__(self, field: Field) -> Field:
@@ -185,12 +188,12 @@ class SeidelAberrations(nn.Module):
         v: The vertical position of the object field point
     """
 
-    coefficients: Union[Array, Callable[[PRNGKey], Array]]
-    f: float
-    n: float
-    NA: float
-    u: float
-    v: float
+    coefficients: ArrayLike | Callable[[PRNGKey], Array]
+    f: ScalarLike
+    n: ScalarLike
+    NA: ScalarLike
+    u: ScalarLike
+    v: ScalarLike
 
     @nn.compact
     def __call__(self, field: Field) -> Field:
@@ -233,11 +236,11 @@ class ZernikeAberrations(nn.Module):
             have same length as coefficients.
     """
 
-    coefficients: Union[Array, Callable[[PRNGKey], Array]]
-    f: float
-    n: float
-    NA: float
-    ansi_indices: Array
+    coefficients: ArrayLike | Callable[[PRNGKey], Array]
+    f: ArrayLike
+    n: ArrayLike
+    NA: ArrayLike
+    ansi_indices: Sequence[int]
 
     @nn.compact
     def __call__(self, field: Field) -> Field:
diff --git a/src/chromatix/elements/propagation.py b/src/chromatix/elements/propagation.py
index 0ac004e..d9089ba 100644
--- a/src/chromatix/elements/propagation.py
+++ b/src/chromatix/elements/propagation.py
@@ -1,10 +1,12 @@
-from typing import Callable, Literal, Optional, Tuple, Union
+from typing import Callable, Literal
 
 import flax.linen as nn
-from chex import Array, PRNGKey
+from chex import PRNGKey
+from jax import Array
 
 from chromatix.elements.utils import Trainable, register
 from chromatix.field import crop, pad
+from chromatix.typing import ArrayLike
 
 from ..field import Field
 from ..functional import (
@@ -81,11 +83,11 @@ class Propagate(nn.Module):
             to True.
     """
 
-    z: Union[float, Array, Callable[[PRNGKey], Array]]
-    n: Union[float, Callable[[PRNGKey], Array]]
+    z: ArrayLike | Callable[[PRNGKey], Array]
+    n: ArrayLike | Callable[[PRNGKey], Array]
     N_pad: int = 0
     cval: float = 0
-    kykx: Union[Array, Tuple[float, float]] = (0.0, 0.0)
+    kykx: ArrayLike | tuple[float, float] = (0.0, 0.0)
     method: Literal["transform", "transfer", "exact", "asm"] = "exact"
     mode: Literal["full", "same"] = "same"
     cache_propagator: bool = True
@@ -128,6 +130,8 @@ def __call__(self, field: Field) -> Field:
                     "kernel",
                     lambda: compute_asm_propagator(*propagator_args),
                 )
+            else:
+                raise NotImplementedError
             field = kernel_propagate(field, propagator.value)
             if self.mode == "same":
                 field = crop(field, self.N_pad)
@@ -207,12 +211,12 @@ class KernelPropagate(nn.Module):
             shape, unlike the functional methods.
     """
 
-    propagator: Union[Array, Callable[[PRNGKey], Array]]
-    z: Optional[Union[float, Array]] = None
-    n: Optional[float] = None
+    propagator: ArrayLike | Callable[[PRNGKey], Array]
+    z: ArrayLike | None = None
+    n: ArrayLike | None = None
     N_pad: int = 0
     cval: float = 0
-    kykx: Union[Array, Tuple[float, float]] = (0.0, 0.0)
+    kykx: ArrayLike | tuple[float, float] = (0.0, 0.0)
     mode: Literal["full", "same"] = "same"
 
     @nn.compact
diff --git a/src/chromatix/elements/samples.py b/src/chromatix/elements/samples.py
index 7cc2418..99670b1 100644
--- a/src/chromatix/elements/samples.py
+++ b/src/chromatix/elements/samples.py
@@ -1,7 +1,10 @@
-from typing import Callable, Union
+from typing import Callable
 
 import flax.linen as nn
-from chex import Array, PRNGKey
+from chex import PRNGKey
+from jax import Array
+
+from chromatix.typing import ArrayLike
 
 from ..field import ScalarField
 from ..functional.samples import thin_sample
@@ -31,9 +34,9 @@ class ThinSample(nn.Module):
             to ``(B... H W 1 1)``
     """
 
-    absorption: Union[Array, Callable[[PRNGKey], Array]]
-    dn: Union[Array, Callable[[PRNGKey], Array]]
-    thickness: Union[Array, Callable[[PRNGKey], Array]]
+    absorption: ArrayLike | Callable[[PRNGKey], Array]
+    dn: ArrayLike | Callable[[PRNGKey], Array]
+    thickness: ArrayLike | Callable[[PRNGKey], Array]
 
     @nn.compact
     def __call__(self, field: ScalarField) -> ScalarField:
diff --git a/src/chromatix/elements/sensors.py b/src/chromatix/elements/sensors.py
index 3a314de..4dea8f5 100644
--- a/src/chromatix/elements/sensors.py
+++ b/src/chromatix/elements/sensors.py
@@ -1,11 +1,12 @@
-from typing import Literal, Optional, Tuple, Union
+from typing import Literal
 
 import flax.linen as nn
 import jax.numpy as jnp
-from chex import Array
-from jax import vmap
+from jax import Array, vmap
 
-from ..field import Field
+from chromatix.typing import ArrayLike, ScalarLike
+
+from ..field import Field, ScalarField, VectorField
 from ..functional import basic_sensor
 from ..ops import init_plane_resample
 
@@ -35,12 +36,12 @@ class BasicSensor(nn.Module):
             name.
     """
 
-    shape: Tuple[int, int]
-    spacing: Union[float, Array]
-    shot_noise_mode: Optional[Literal["approximate", "poisson"]] = None
-    resampling_method: Optional[str] = "linear"
-    reduce_axis: Optional[int] = None
-    reduce_parallel_axis_name: Optional[str] = None
+    shape: tuple[int, int]
+    spacing: ArrayLike
+    shot_noise_mode: Literal["approximate", "poisson"] | None = None
+    resampling_method: str | None = "linear"
+    reduce_axis: int | None = None
+    reduce_parallel_axis_name: str | None = None
 
     def setup(self):
         if self.resampling_method is not None:
@@ -50,8 +51,8 @@ def setup(self):
 
     def __call__(
         self,
-        sensor_input: Union[Field, Array],
-        input_spacing: Optional[Union[float, Array]] = None,
+        sensor_input: Array | Field,
+        input_spacing: ScalarLike | None = None,
         resample: bool = True,
     ) -> Array:
         """
@@ -65,11 +66,15 @@ def __call__(
             resample: Whether to perform resampling or not. Only matters if
                 ``resampling_method`` is ``None``. Defaults to ``True``.
         """
-        if isinstance(sensor_input, Field):
+
+        if (input_spacing is not None) and isinstance(
+            sensor_input, (ScalarField, VectorField)
+        ):
             # WARNING(dd): @copypaste(Microscope) Assumes that field has same
             # spacing at all wavelengths when calculating intensity!
             input_spacing = sensor_input.dx[..., 0, 0].squeeze()
-        input_spacing = jnp.atleast_1d(input_spacing)
+        elif input_spacing is not None:
+            input_spacing = jnp.atleast_1d(input_spacing)
         if resample and self.resampling_method is not None:
             resample_fn = self.resample_fn
         else:
@@ -88,7 +93,7 @@ def __call__(
             noise_key=noise_key,
         )
 
-    def resample(self, resample_input: Array, input_spacing: float) -> Array:
+    def resample(self, resample_input: Array, input_spacing: ScalarLike) -> Array:
         """
         Resample the given ``resample_input`` to the pixels of the sensor.
 
diff --git a/src/chromatix/elements/sources.py b/src/chromatix/elements/sources.py
index 3a00275..3202c09 100644
--- a/src/chromatix/elements/sources.py
+++ b/src/chromatix/elements/sources.py
@@ -1,18 +1,24 @@
-from typing import Callable, Optional, Tuple, Union
-import numpy as np
+from typing import Callable
+
 import flax.linen as nn
-from chex import Array, PRNGKey
+import numpy as np
+from chex import PRNGKey
+from jax import Array
+
 from chromatix.elements.utils import register
-from chromatix.field import Field
+from chromatix.field import Field, ScalarField, VectorField
 from chromatix.functional.sources import (
     generic_field,
     objective_point_source,
     plane_wave,
     point_source,
 )
+from chromatix.typing import ArrayLike, ScalarLike
 
 __all__ = ["PointSource", "ObjectivePointSource", "PlaneWave", "GenericField"]
 
+FieldPupil = Callable[[Field], Field]
+
 
 class PointSource(nn.Module):
     """
@@ -42,20 +48,20 @@ class PointSource(nn.Module):
         epsilon: Value added to denominators for numerical stability.
     """
 
-    shape: Tuple[int, int]
-    dx: Union[float, Array]
-    spectrum: Union[float, Array]
-    spectral_density: Union[float, Array]
-    z: Union[float, Callable[[PRNGKey], float]]
-    n: Union[float, Callable[[PRNGKey], float]]
-    power: Union[float, Callable[[PRNGKey], float]] = 1.0
-    amplitude: Union[float, Array, Callable[[PRNGKey], Array]] = 1.0
-    pupil: Optional[Callable[[Field], Field]] = None
+    shape: tuple[int, int]
+    dx: ScalarLike
+    spectrum: ScalarLike
+    spectral_density: ScalarLike
+    z: ScalarLike | Callable[[PRNGKey], Array]
+    n: ScalarLike | Callable[[PRNGKey], Array]
+    power: ScalarLike | Callable[[PRNGKey], Array] = 1.0
+    amplitude: ScalarLike | Callable[[PRNGKey], Array] = 1.0
+    pupil: FieldPupil | None = None
     scalar: bool = True
-    epsilon: float = np.finfo(np.float32).eps,
+    epsilon: float = float(np.finfo(np.float32).eps)
 
     @nn.compact
-    def __call__(self) -> Field:
+    def __call__(self) -> ScalarField | VectorField:
         power = register(self, "power")
         z = register(self, "z")
         n = register(self, "n")
@@ -71,7 +77,7 @@ def __call__(self) -> Field:
             amplitude,
             self.pupil,
             self.scalar,
-            self.epsilon
+            self.epsilon,
         )
 
 
@@ -103,21 +109,20 @@ class ObjectivePointSource(nn.Module):
         scalar: Whether the result should be ``ScalarField`` (if True) or
             ``VectorField`` (if False). Defaults to True.
     """
-
-    shape: Tuple[int, int]
-    dx: Union[float, Array]
-    spectrum: Union[float, Array]
-    spectral_density: Union[float, Array]
-    f: Union[float, Callable[[PRNGKey], float]]
-    n: Union[float, Callable[[PRNGKey], float]]
-    NA: Union[float, Callable[[PRNGKey], float]]
-    power: Union[float, Callable[[PRNGKey], float]] = 1.0
-    amplitude: Union[float, Array, Callable[[PRNGKey], Array]] = 1.0
-    offset: Union[Array, Tuple[float, float]] = (0.0, 0.0)
+    shape: tuple[int, int]
+    dx: ScalarLike
+    spectrum: ScalarLike
+    spectral_density: ScalarLike
+    f: ScalarLike | Callable[[PRNGKey], Array]
+    n: ScalarLike | Callable[[PRNGKey], Array]
+    NA: ScalarLike | Callable[[PRNGKey], Array]
+    power: ScalarLike | Callable[[PRNGKey], Array] = 1.0
+    amplitude: ScalarLike | Callable[[PRNGKey], Array] = 1.0
+    offset: ArrayLike | tuple[float, float] = (0.0, 0.0)
     scalar: bool = True
 
     @nn.compact
-    def __call__(self, z: float) -> Field:
+    def __call__(self, z: ScalarLike) -> ScalarField | VectorField:
         f = register(self, "f")
         n = register(self, "n")
         NA = register(self, "NA")
@@ -169,18 +174,18 @@ class PlaneWave(nn.Module):
             ``VectorField`` (if False). Defaults to True.
     """
 
-    shape: Tuple[int, int]
-    dx: Union[float, Array]
-    spectrum: Union[float, Array]
-    spectral_density: Union[float, Array]
-    power: Union[float, Callable[[PRNGKey], float]] = 1.0
-    amplitude: Union[float, Array, Callable[[PRNGKey], Array]] = 1.0
-    kykx: Union[Array, Tuple[float, float]] = (0.0, 0.0)
-    pupil: Optional[Callable[[Field], Field]] = None
+    shape: tuple[int, int]
+    dx: ScalarLike
+    spectrum: ScalarLike
+    spectral_density: ScalarLike
+    power: ScalarLike | Callable[[PRNGKey], Array] = 1.0
+    amplitude: ScalarLike | Callable[[PRNGKey], Array] = 1.0
+    kykx: ArrayLike | tuple[float, float] = (0.0, 0.0)
+    pupil: FieldPupil | None = None
     scalar: bool = True
 
     @nn.compact
-    def __call__(self) -> Field:
+    def __call__(self) -> ScalarField | VectorField:
         kykx = register(self, "kykx")
         power = register(self, "power")
         amplitude = register(self, "amplitude")
@@ -220,17 +225,17 @@ class GenericField(nn.Module):
             ``VectorField`` (if False). Defaults to True.
     """
 
-    dx: Union[float, Array]
-    spectrum: Union[float, Array]
-    spectral_density: Union[float, Array]
-    amplitude: Union[Array, Callable[[PRNGKey], Array]]
-    phase: Union[Array, Callable[[PRNGKey], Array]]
-    power: Union[float, Callable[[PRNGKey], float]] = 1.0
-    pupil: Optional[Callable[[Field], Field]] = None
+    dx: ScalarLike
+    spectrum: ScalarLike
+    spectral_density: ScalarLike
+    amplitude: ArrayLike | Callable[[PRNGKey], Array]
+    phase: ArrayLike | Callable[[PRNGKey], Array]
+    power: ScalarLike | Callable[[PRNGKey], Array] = 1.0
+    pupil: FieldPupil | None = None
     scalar: bool = True
 
     @nn.compact
-    def __call__(self) -> Field:
+    def __call__(self) -> ScalarField | VectorField:
         amplitude = register(self, "amplitude")
         phase = register(self, "phase")
         power = register(self, "power")
diff --git a/src/chromatix/elements/utils.py b/src/chromatix/elements/utils.py
index 0917cd4..47e0897 100644
--- a/src/chromatix/elements/utils.py
+++ b/src/chromatix/elements/utils.py
@@ -1,8 +1,9 @@
 from dataclasses import dataclass
 from typing import Any, Callable
 
-from chex import Array, PRNGKey
+from chex import PRNGKey
 from flax import linen as nn
+from jax import Array
 
 
 @dataclass
@@ -135,17 +136,12 @@ def register(
     try:
         init = getattr(module, name)
     except AttributeError:
-        print("Variable does not exist.")
+        raise AttributeError
 
     if isinstance(init, Trainable):
         return module.param(f"_{name}", parse_init(init.val), *args)
     else:
-        return module.variable(
-            "state",
-            f"_{name}",
-            parse_init(init),
-            *args,
-        ).value
+        return module.variable("state", f"_{name}", parse_init(init), *args).value
 
 
 def parse_init(x: Any) -> Callable:
diff --git a/src/chromatix/field.py b/src/chromatix/field.py
index fc79622..7b9d1ae 100644
--- a/src/chromatix/field.py
+++ b/src/chromatix/field.py
@@ -1,12 +1,16 @@
 from __future__ import annotations
 
 from numbers import Number
-from typing import Any, Optional, Tuple, Union
+from typing import Any, TypeVar
 
 import jax.numpy as jnp
-from chex import Array, assert_equal_shape, assert_rank
+from chex import assert_equal_shape, assert_rank
 from einops import rearrange
 from flax import struct
+from jax import Array
+from typing_extensions import Self
+
+from chromatix.typing import ArrayLike, ScalarLike
 
 from .utils.shapes import (
     _broadcast_1d_to_channels,
@@ -15,7 +19,7 @@
 )
 
 
-class Field(struct.PyTreeNode):
+class BaseField(struct.PyTreeNode):
     """
     A container that describes the chromatic light field at a 2D plane.
 
@@ -201,17 +205,17 @@ def power(self) -> Array:
         return jnp.sum(self.intensity, axis=(-4, -3), keepdims=True) * area
 
     @property
-    def shape(self) -> Tuple[int, ...]:
+    def shape(self) -> tuple[int, ...]:
         """Shape of the complex field."""
         return self.u.shape
 
     @property
-    def spatial_shape(self) -> Tuple[int, int]:
+    def spatial_shape(self) -> tuple[int, int]:
         """Only the height and width of the complex field."""
-        return self.u.shape[self.spatial_dims[0] : self.spatial_dims[1] + 1]
+        return (self.u.shape[self.spatial_dims[0]], self.u.shape[self.spatial_dims[1]])
 
     @property
-    def spatial_dims(self) -> Tuple[int, int]:
+    def spatial_dims(self) -> tuple[int, int]:
         """Dimensions representing the height and width of the complex field."""
         return (-4, -3)
 
@@ -221,93 +225,93 @@ def ndim(self) -> int:
         return self.u.ndim
 
     @property
-    def conj(self) -> Array:
+    def conj(self) -> Self:
         """conjugate of the complex field, as a field of the same shape."""
         return self.replace(u=jnp.conj(self.u))
 
-    def __add__(self, other: Union[Number, jnp.ndarray, Field]) -> Field:
+    def __add__(self, other: ScalarLike | ArrayLike | Field) -> Self:
         if isinstance(other, jnp.ndarray) or isinstance(other, Number):
             return self.replace(u=self.u + other)
-        elif isinstance(other, Field):
+        elif isinstance(other, (ScalarField, VectorField)):
             return self.replace(u=self.u + other.u)
         else:
             return NotImplemented
 
-    def __radd__(self, other: Any) -> Field:
+    def __radd__(self, other: Any) -> Self:
         return self + other
 
-    def __sub__(self, other: Union[Number, jnp.ndarray, Field]) -> Field:
+    def __sub__(self, other: ScalarLike | ArrayLike | Field) -> Self:
         if isinstance(other, jnp.ndarray) or isinstance(other, Number):
             return self.replace(u=self.u - other)
-        elif isinstance(other, Field):
+        elif isinstance(other, (ScalarField, VectorField)):
             return self.replace(u=self.u - other.u)
         else:
             return NotImplemented
 
-    def __rsub__(self, other: Any) -> Field:
+    def __rsub__(self, other: Any) -> Self:
         return (-1 * self) + other
 
-    def __mul__(self, other: Union[Number, jnp.ndarray, Field]) -> Field:
+    def __mul__(self, other: ScalarLike | ArrayLike | Field) -> Self:
         if isinstance(other, jnp.ndarray) or isinstance(other, Number):
             return self.replace(u=self.u * other)
-        elif isinstance(other, Field):
+        elif isinstance(other, (ScalarField, VectorField)):
             return self.replace(u=self.u * other.u)
         else:
             return NotImplemented
 
-    def __rmul__(self, other: Any) -> Field:
+    def __rmul__(self, other: Any) -> Self:
         return self * other
 
-    def __matmul__(self, other: jnp.array) -> Field:
+    def __matmul__(self, other: ArrayLike) -> Self:
         return self.replace(u=jnp.matmul(self.u, other))
 
-    def __rmatmul__(self, other: jnp.array) -> Field:
+    def __rmatmul__(self, other: ArrayLike) -> Self:
         return self.replace(u=jnp.matmul(other, self.u.squeeze()))
 
-    def __truediv__(self, other: Union[Number, jnp.ndarray, Field]) -> Field:
+    def __truediv__(self, other: ScalarLike | ArrayLike | Field) -> Self:
         if isinstance(other, jnp.ndarray) or isinstance(other, Number):
             return self.replace(u=self.u / other)
-        elif isinstance(other, Field):
+        elif isinstance(other, (ScalarField, VectorField)):
             return self.replace(u=self.u / other.u)
         else:
             return NotImplemented
 
-    def __rtruediv__(self, other: Any) -> Field:
+    def __rtruediv__(self, other: Any) -> Self:
         return self.replace(u=other / self.u)
 
-    def __floordiv__(self, other: Union[Number, jnp.ndarray, Field]) -> Field:
+    def __floordiv__(self, other: ScalarLike | ArrayLike | Field) -> Self:
         if isinstance(other, jnp.ndarray) or isinstance(other, Number):
             return self.replace(u=self.u // other)
-        elif isinstance(other, Field):
+        elif isinstance(other, (ScalarField, VectorField)):
             return self.replace(u=self.u // other.u)
         else:
             return NotImplemented
 
-    def __rfloordiv__(self, other: Any) -> Field:
+    def __rfloordiv__(self, other: Any) -> Self:
         return self.replace(u=other // self.u)
 
-    def __mod__(self, other: Union[Number, jnp.ndarray, Field]) -> Field:
+    def __mod__(self, other: ScalarLike | ArrayLike | Field) -> Self:
         if isinstance(other, jnp.ndarray) or isinstance(other, Number):
             return self.replace(u=self.u % other)
-        elif isinstance(other, Field):
+        elif isinstance(other, (ScalarField, VectorField)):
             return self.replace(u=self.u % other.u)
         else:
             return NotImplemented
 
-    def __rmod__(self, other: Any) -> Field:
+    def __rmod__(self, other: Any) -> Self:
         return self.replace(u=other % self.u)
 
 
-class ScalarField(Field):
+class ScalarField(BaseField):
     @classmethod
     def create(
         cls,
-        dx: Union[float, Array],
-        spectrum: Union[float, Array],
-        spectral_density: Union[float, Array],
-        u: Optional[Array] = None,
-        shape: Optional[Tuple[int, int]] = None,
-    ) -> Field:
+        dx: ScalarLike,
+        spectrum: ScalarLike,
+        spectral_density: ScalarLike,
+        u: Array | None = None,
+        shape: tuple[int, int] | None = None,
+    ) -> Self:
         """
         Create a scalar approximation ``Field`` object in a convenient way.
 
@@ -340,35 +344,41 @@ def create(
                 if ``u`` is provided. If ``u`` is not provided, then ``shape``
                 must be provided.
         """
-        dx: Array = jnp.atleast_1d(dx)
-        spectrum: Array = jnp.atleast_1d(spectrum)
-        spectral_density: Array = jnp.atleast_1d(spectral_density)
+        # Parsing dx
+        _dx = jnp.atleast_1d(dx)
+        if _dx.ndim == 1:
+            _dx = jnp.stack([_dx, _dx])
+        assert_rank(_dx, 2)  # dx should have shape (2, C) here
+
+        # Parsing spectrum and density
+        _spectrum = jnp.atleast_1d(spectrum)
+        _spectral_density = jnp.atleast_1d(spectral_density)
+        assert_equal_shape([_spectrum, _spectral_density])
+        _spectral_density = _spectral_density / jnp.sum(_spectral_density)
+
+        # Parsing u
         if u is None:
             assert shape is not None, "Must specify shape if u is None"
-            u = jnp.empty((1, *shape, spectrum.size, 1), dtype=jnp.complex64)
-        ndim = len(u.shape)
-        assert (
-            ndim >= 5
-        ), "Field must be Array with at least 5 dimensions: (B... H W C 1)."
-        assert u.shape[-1] == 1, "Last dimension must be 1 for scalar fields."
-        assert_equal_shape([spectrum, spectral_density])
-        spectral_density = spectral_density / jnp.sum(spectral_density)
-        if dx.ndim == 1:
-            dx = jnp.stack([dx, dx])
-        assert_rank(dx, 2)  # dx should have shape (2, C) here
-        return cls(u, dx, spectrum, spectral_density)
-
-
-class VectorField(Field):
+            _u = jnp.empty((1, *shape, _spectrum.size, 1), dtype=jnp.complex64)
+        else:
+            _u = jnp.array(u)
+        ndim = len(_u.shape)
+        assert ndim >= 5, "Field must be Array with at least 5 dimensions: (B H W C 3)."
+        assert _u.shape[-1] == 1, "Last dimension must be 1 for scalar fields."
+
+        return cls(_u, _dx, _spectrum, _spectral_density)
+
+
+class VectorField(BaseField):
     @classmethod
     def create(
         cls,
-        dx: Union[float, Array],
-        spectrum: Union[float, Array],
-        spectral_density: Union[float, Array],
-        u: Optional[Array] = None,
-        shape: Optional[Tuple[int, int]] = None,
-    ) -> Field:
+        dx: ScalarLike,
+        spectrum: ScalarLike,
+        spectral_density: ScalarLike,
+        u: ArrayLike | None = None,
+        shape: tuple[int, int] | None = None,
+    ) -> Self:
         """
         Create a vectorial ``Field`` object in a convenient way.
 
@@ -401,23 +411,30 @@ def create(
                 if ``u`` is provided. If ``u`` is not provided, then ``shape``
                 must be provided.
         """
-        dx: Array = jnp.atleast_1d(dx)
-        spectrum: Array = jnp.atleast_1d(spectrum)
-        spectral_density: Array = jnp.atleast_1d(spectral_density)
+
+        # Parsing dx
+        _dx = jnp.atleast_1d(dx)
+        if _dx.ndim == 1:
+            _dx = jnp.stack([_dx, _dx])
+        assert_rank(_dx, 2)  # dx should have shape (2, C) here
+
+        # Parsing spectrum and density
+        _spectrum = jnp.atleast_1d(spectrum)
+        _spectral_density = jnp.atleast_1d(spectral_density)
+        assert_equal_shape([_spectrum, _spectral_density])
+        _spectral_density = _spectral_density / jnp.sum(_spectral_density)
+
+        # Parsing u
         if u is None:
             assert shape is not None, "Must specify shape if u is None"
-            u = jnp.empty((1, *shape, spectrum.size, 3), dtype=jnp.complex64)
-        ndim = len(u.shape)
-        assert (
-            ndim >= 5
-        ), "Field must be Array with at least 5 dimensions: (B... H W C 3)."
-        assert u.shape[-1] == 3, "Last dimension must be 3 for vectorial fields."
-        assert_equal_shape([spectrum, spectral_density])
-        spectral_density = spectral_density / jnp.sum(spectral_density)
-        if dx.ndim == 1:
-            dx = jnp.stack([dx, dx])
-        assert_rank(dx, 2)  # dx should have shape (2, C) here
-        return cls(u, dx, spectrum, spectral_density)
+            _u = jnp.empty((1, *shape, _spectrum.size, 3), dtype=jnp.complex64)
+        else:
+            _u = jnp.array(u)
+        ndim = len(_u.shape)
+        assert ndim >= 5, "Field must be Array with at least 5 dimensions: (B H W C 3)."
+        assert _u.shape[-1] == 3, "Last dimension must be 3 for vectorial fields."
+
+        return cls(_u, _dx, _spectrum, _spectral_density)
 
     @property
     def jones_vector(self) -> Array:
@@ -427,7 +444,10 @@ def jones_vector(self) -> Array:
         return self.u / norm
 
 
-def pad(field: Field, pad_width: Union[int, Tuple[int, int]], cval: float = 0) -> Field:
+Field = TypeVar("Field", ScalarField, VectorField)
+
+
+def pad(field: Field, pad_width: int | tuple[int, int], cval: float = 0) -> Field:
     """
     Pad the `field` with zeros in one or two dimensions.
     Args:
@@ -445,7 +465,7 @@ def pad(field: Field, pad_width: Union[int, Tuple[int, int]], cval: float = 0) -
     return field.replace(u=u)
 
 
-def crop(field: Field, crop_width: Union[int, Tuple[int, int]]) -> Field:
+def crop(field: Field, crop_width: int | tuple[int, int]) -> Field:
     """
     Crop the `field` by removing pixels from the edges.
     Args:
@@ -461,7 +481,7 @@ def crop(field: Field, crop_width: Union[int, Tuple[int, int]]) -> Field:
     return field.replace(u=field.u[tuple(crop)])
 
 
-def shift(field: Field, shiftby: Union[int, Tuple[int, int]]) -> Field:
+def shift(field: Field, shiftby: int | tuple[int, int]) -> Field:
     """
     Shift the `field` by an integer number of pixels in one or two dimensions.
     Args:
diff --git a/src/chromatix/functional/amplitude_masks.py b/src/chromatix/functional/amplitude_masks.py
index 518409f..ca956cd 100644
--- a/src/chromatix/functional/amplitude_masks.py
+++ b/src/chromatix/functional/amplitude_masks.py
@@ -1,7 +1,9 @@
 import jax.numpy as jnp
-from chex import Array, assert_rank
+from chex import assert_rank
+from jax import Array
+
+from chromatix import Field
 
-from ..field import Field
 from ..utils.shapes import _broadcast_2d_to_spatial
 
 __all__ = ["amplitude_change"]
diff --git a/src/chromatix/functional/convenience.py b/src/chromatix/functional/convenience.py
index 1424a63..ec379ea 100644
--- a/src/chromatix/functional/convenience.py
+++ b/src/chromatix/functional/convenience.py
@@ -1,11 +1,12 @@
 import jax.numpy as jnp
 
-from chromatix import Field
+from chromatix.field import Field
+from chromatix.typing import ScalarLike
 from chromatix.utils import _squeeze_grid_to_2d
 from chromatix.utils.fft import fft
 
 
-def optical_fft(field: Field, z: float, n: float) -> Field:
+def optical_fft(field: Field, z: ScalarLike, n: ScalarLike) -> Field:
     """
     Computes the optical ``fft`` or ``ifft`` on an incoming ``Field`` propagated
     by ``z``, depending on the sign of ``z`` (which is a scalar value that may
diff --git a/src/chromatix/functional/lenses.py b/src/chromatix/functional/lenses.py
index 631bd6c..286c429 100644
--- a/src/chromatix/functional/lenses.py
+++ b/src/chromatix/functional/lenses.py
@@ -1,17 +1,18 @@
-from typing import Optional
-
 import jax.numpy as jnp
 
+from chromatix import Field
 from chromatix.functional.convenience import optical_fft
+from chromatix.typing import ScalarLike
 
-from ..field import Field
 from ..utils import l2_sq_norm
 from .pupils import circular_pupil
 
 __all__ = ["thin_lens", "ff_lens", "df_lens"]
 
 
-def thin_lens(field: Field, f: float, n: float, NA: Optional[float] = None) -> Field:
+def thin_lens(
+    field: Field, f: ScalarLike, n: ScalarLike, NA: ScalarLike | None = None
+) -> Field:
     """
     Applies a thin lens placed directly after the incoming ``Field``.
 
@@ -37,9 +38,9 @@ def thin_lens(field: Field, f: float, n: float, NA: Optional[float] = None) -> F
 
 def ff_lens(
     field: Field,
-    f: float,
-    n: float,
-    NA: Optional[float] = None,
+    f: ScalarLike,
+    n: ScalarLike,
+    NA: ScalarLike | None = None,
     inverse: bool = False,
 ) -> Field:
     """
@@ -67,10 +68,10 @@ def ff_lens(
 
 def df_lens(
     field: Field,
-    d: float,
-    f: float,
-    n: float,
-    NA: Optional[float] = None,
+    d: ScalarLike,
+    f: ScalarLike,
+    n: ScalarLike,
+    NA: ScalarLike | None = None,
     inverse: bool = False,
 ) -> Field:
     """
diff --git a/src/chromatix/functional/phase_masks.py b/src/chromatix/functional/phase_masks.py
index be06325..f8d9a46 100644
--- a/src/chromatix/functional/phase_masks.py
+++ b/src/chromatix/functional/phase_masks.py
@@ -1,9 +1,12 @@
-from typing import Tuple
 import jax
 import jax.numpy as jnp
-from chex import Array, assert_rank
-from chromatix.field import Field
-from chromatix.utils.shapes import _broadcast_2d_to_spatial
+from chex import assert_rank
+from jax import Array
+
+from chromatix.typing import ArrayLike
+
+from ..field import Field, ScalarField, VectorField
+from ..utils.shapes import _broadcast_2d_to_spatial
 
 __all__ = ["phase_change", "wrap_phase", "spectrally_modulate_phase"]
 
@@ -30,7 +33,9 @@ def phase_change(field: Field, phase: Array, spectrally_modulate: bool = True) -
 
 
 @jax.custom_jvp
-def wrap_phase(phase: Array, limits: Tuple[float, float] = (-jnp.pi, jnp.pi)) -> Array:
+def wrap_phase(
+    phase: ArrayLike, limits: ArrayLike | tuple[float, float] = (-jnp.pi, jnp.pi)
+) -> Array:
     """
     Wraps values of ``phase`` to the range given by ``limits``.
 
@@ -54,11 +59,11 @@ def wrap_phase(phase: Array, limits: Tuple[float, float] = (-jnp.pi, jnp.pi)) ->
 
 
 @wrap_phase.defjvp
-def wrap_phase_jvp(primals: Tuple, tangents: Tuple) -> Tuple:
+def wrap_phase_jvp(primals: tuple, tangents: tuple) -> tuple:
     return wrap_phase(*primals), tangents[0]
 
 
-def spectrally_modulate_phase(phase: Array, field: Field) -> Array:
+def spectrally_modulate_phase(phase: Array, field: ScalarField | VectorField) -> Array:
     """Spectrally modulates a given ``phase`` for multiple wavelengths."""
     central_wavelength = field.spectrum[..., 0, 0].squeeze()
     spectral_modulation = central_wavelength / field.spectrum
diff --git a/src/chromatix/functional/polarizers.py b/src/chromatix/functional/polarizers.py
index efb692e..c9de7fe 100644
--- a/src/chromatix/functional/polarizers.py
+++ b/src/chromatix/functional/polarizers.py
@@ -1,8 +1,7 @@
-from typing import Union
-
 import jax.numpy as jnp
-from chex import Array
+from jax import Array
 
+from chromatix.typing import ComplexScalarLike, ScalarLike
 from chromatix.utils.utils import matvec
 
 from ..field import VectorField
@@ -28,7 +27,7 @@
 ]
 
 
-def jones_vector(theta: float, beta: float) -> Array:
+def jones_vector(theta: ScalarLike, beta: ScalarLike) -> Array:
     """Generates a Jones vector given by [cos(theta), sin(theta)exp(1j*beta)].
 
     Args:
@@ -46,7 +45,7 @@ def jones_vector(theta: float, beta: float) -> Array:
     )
 
 
-def linear(theta: float) -> Array:
+def linear(theta: ScalarLike) -> Array:
     """Generates a Jones vector for linearly polarized
     light with an angle $\theta$ w.r.t. to the horizontal.
 
@@ -81,10 +80,10 @@ def right_circular() -> Array:
 
 def polarizer(
     field: VectorField,
-    J00: Union[float, complex, Array],
-    J01: Union[float, complex, Array],
-    J10: Union[float, complex, Array],
-    J11: Union[float, complex, Array],
+    J00: ComplexScalarLike,
+    J01: ComplexScalarLike,
+    J10: ComplexScalarLike,
+    J11: ComplexScalarLike,
 ) -> VectorField:
     """Applies a Jones matrix with given components to the field.
     Note that the components here refer to the common choice of coordinate
@@ -106,7 +105,7 @@ def polarizer(
     return field.replace(u=matvec(LP, field.u))
 
 
-def linear_polarizer(field: VectorField, angle: float) -> VectorField:
+def linear_polarizer(field: VectorField, angle: ScalarLike) -> VectorField:
     """Applies a linear polarizer with a given angle to the incoming field.
 
     Args:
@@ -159,7 +158,7 @@ def right_circular_polarizer(field: VectorField) -> VectorField:
 
 
 def phase_retarder(
-    field: VectorField, theta: float, eta: float, phi: float
+    field: VectorField, theta: ScalarLike, eta: ScalarLike, phi: ScalarLike
 ) -> VectorField:
     """Applies a general purpose retardation matrix with angle w.r.t horizontal theta,
     relative phase change eta and circularity phi.
@@ -182,7 +181,7 @@ def phase_retarder(
     return polarizer(field, J00, J01, J10, J11)
 
 
-def wave_plate(field: VectorField, theta: float, eta: float) -> VectorField:
+def wave_plate(field: VectorField, theta: ScalarLike, eta: ScalarLike) -> VectorField:
     """Applies a general waveplate with angle theta and delay eta to the field.
 
     Args:
@@ -196,7 +195,7 @@ def wave_plate(field: VectorField, theta: float, eta: float) -> VectorField:
     return phase_retarder(field, theta, eta, phi=0)
 
 
-def halfwave_plate(field: VectorField, theta: float) -> VectorField:
+def halfwave_plate(field: VectorField, theta: ScalarLike) -> VectorField:
     """Applies a halfwave plate with angle theta to the incoming field.
 
     Args:
@@ -209,7 +208,7 @@ def halfwave_plate(field: VectorField, theta: float) -> VectorField:
     return phase_retarder(field, theta, eta=jnp.pi, phi=0)
 
 
-def quarterwave_plate(field: VectorField, theta: float) -> VectorField:
+def quarterwave_plate(field: VectorField, theta: ScalarLike) -> VectorField:
     """Applies a quarterwave plate with angle theta to the incoming field.
 
     Args:
@@ -222,7 +221,9 @@ def quarterwave_plate(field: VectorField, theta: float) -> VectorField:
     return phase_retarder(field, theta, eta=jnp.pi / 2, phi=0)
 
 
-def universal_compensator(field: VectorField, retA: float, retB: float) -> VectorField:
+def universal_compensator(
+    field: VectorField, retA: ScalarLike, retB: ScalarLike
+) -> VectorField:
     """Applies the Universal Polarizer for the LC-PolScope to the incoming field.
 
     Args:
diff --git a/src/chromatix/functional/propagation.py b/src/chromatix/functional/propagation.py
index 82b0a43..c9c7b6c 100644
--- a/src/chromatix/functional/propagation.py
+++ b/src/chromatix/functional/propagation.py
@@ -1,15 +1,16 @@
-from typing import Literal, Tuple, Union
+from typing import Literal
 
 import jax
 import jax.numpy as jnp
 import numpy as np
-from chex import Array
+from jax import Array
 
 from chromatix.field import crop, pad
 from chromatix.functional.convenience import optical_fft
+from chromatix.typing import ArrayLike, ScalarLike
 from chromatix.utils.fft import fft, ifft
 
-from ..field import Field
+from ..field import Field, ScalarField, VectorField
 from ..utils import _broadcast_1d_to_grid, _broadcast_1d_to_innermost_batch, l2_sq_norm
 
 __all__ = [
@@ -31,9 +32,9 @@
 
 def transform_propagate(
     field: Field,
-    z: float,
-    n: float,
-    N_pad: Union[int, Tuple[int, int]],
+    z: ScalarLike,
+    n: ScalarLike,
+    N_pad: int | tuple[int, int],
     cval: float = 0,
     skip_initial_phase: bool = False,
     skip_final_phase: bool = False,
@@ -81,10 +82,10 @@ def transform_propagate(
 
 
 def compute_sas_precompensation(
-    field: Field,
-    z: float,
-    n: float,
-) -> Field:
+    field: ScalarField | VectorField,
+    z: ScalarLike,
+    n: ScalarLike,
+) -> Array:
     sz = np.array(field.spatial_shape)
     kz = 2 * z * jnp.pi * n / field.spectrum
     s = field.spectrum * field.k_grid / n
@@ -105,8 +106,8 @@ def compute_sas_precompensation(
 
 def transform_propagate_sas(
     field: Field,
-    z: float,
-    n: float,
+    z: ScalarLike,
+    n: ScalarLike,
     cval: float = 0,
     skip_initial_phase: bool = False,
     skip_final_phase: bool = False,
@@ -141,10 +142,10 @@ def transform_propagate_sas(
     # Don't change this pad_factor, only 2 is supported
     pad_factor = 2
     sz = np.array(field.spatial_shape)
-    N_pad = sz // pad_factor
+    N_pad = tuple(sz // pad_factor)
     field = pad(field, N_pad, cval=cval)
 
-    def _forward(field: Field, z) -> Tuple[Array, Array]:
+    def _forward(field: Field, z) -> tuple[Array, Array]:
         delta_H = compute_sas_precompensation(field, z, n)
         field = kernel_propagate(field, delta_H)
         field = transform_propagate(
@@ -152,7 +153,7 @@ def _forward(field: Field, z) -> Tuple[Array, Array]:
         )
         return field.u, field._dx
 
-    def _inverse(field: Field, z) -> Tuple[Array, Array]:
+    def _inverse(field: Field, z) -> tuple[Array, Array]:
         field = transform_propagate(
             field, z, n, 0, 0, skip_initial_phase, skip_final_phase
         )
@@ -167,11 +168,11 @@ def _inverse(field: Field, z) -> Tuple[Array, Array]:
 
 def transfer_propagate(
     field: Field,
-    z: Union[float, Array],
-    n: float,
+    z: ScalarLike,
+    n: ScalarLike,
     N_pad: int,
     cval: float = 0,
-    kykx: Union[Array, Tuple[float, float]] = (0.0, 0.0),
+    kykx: ArrayLike | tuple[float, float] = (0.0, 0.0),
     mode: Literal["full", "same"] = "full",
 ) -> Field:
     """
@@ -206,11 +207,11 @@ def transfer_propagate(
 
 def exact_propagate(
     field: Field,
-    z: Union[float, Array],
-    n: float,
+    z: ScalarLike,
+    n: ScalarLike,
     N_pad: int,
     cval: float = 0,
-    kykx: Union[Array, Tuple[float, float]] = (0.0, 0.0),
+    kykx: ArrayLike | tuple[float, float] = (0.0, 0.0),
     mode: Literal["full", "same"] = "full",
 ) -> Field:
     """
@@ -246,13 +247,13 @@ def exact_propagate(
 
 def asm_propagate(
     field: Field,
-    z: Union[float, Array],
-    n: float,
+    z: ScalarLike,
+    n: ScalarLike,
     N_pad: int,
     cval: float = 0,
-    kykx: Union[Array, Tuple[float, float]] = (0.0, 0.0),
+    kykx: ArrayLike | tuple[float, float] = (0.0, 0.0),
     bandlimit: bool = False,
-    shift_yx: Union[Array, Tuple[float, float]] = (0.0, 0.0),
+    shift_yx: ArrayLike | tuple[float, float] = (0.0, 0.0),
     mode: Literal["full", "same"] = "full",
 ) -> Field:
     """
@@ -292,7 +293,7 @@ def asm_propagate(
     return field
 
 
-def kernel_propagate(field: Field, propagator: Array) -> Field:
+def kernel_propagate(field: Field, propagator: ArrayLike) -> Field:
     """
     Propagate an incoming ``Field`` by the given propagation kernel
     (``propagator``). This amounts to performing a Fourier convolution of the
@@ -304,10 +305,10 @@ def kernel_propagate(field: Field, propagator: Array) -> Field:
 
 
 def compute_transfer_propagator(
-    field: Field,
-    z: Union[float, Array],
-    n: float,
-    kykx: Union[Array, Tuple[float, float]] = (0.0, 0.0),
+    field: ScalarField | VectorField,
+    z: ScalarLike,
+    n: ScalarLike,
+    kykx: ArrayLike | tuple[float, float] = (0.0, 0.0),
 ) -> Array:
     """
     Compute propagation kernel for Fresnel propagation.
@@ -328,10 +329,10 @@ def compute_transfer_propagator(
 
 
 def compute_exact_propagator(
-    field: Field,
-    z: Union[float, Array],
-    n: float,
-    kykx: Union[Array, Tuple[float, float]] = (0.0, 0.0),
+    field: ScalarField | VectorField,
+    z: ScalarLike,
+    n: ScalarLike,
+    kykx: ArrayLike | tuple[float, float] = (0.0, 0.0),
 ) -> Array:
     """
     Compute propagation kernel for propagation with no Fresnel approximation.
@@ -358,12 +359,12 @@ def compute_exact_propagator(
 
 
 def compute_asm_propagator(
-    field: Field,
-    z: Union[float, Array],
-    n: float,
-    kykx: Union[Array, Tuple[float, float]] = (0.0, 0.0),
+    field: ScalarField | VectorField,
+    z: ScalarLike,
+    n: ScalarLike,
+    kykx: ArrayLike | tuple[float, float] = (0.0, 0.0),
     bandlimit: bool = False,
-    shift_yx: Union[Array, Tuple[float, float]] = (0.0, 0.0),
+    shift_yx: ArrayLike | tuple[float, float] = (0.0, 0.0),
 ) -> Array:
     """
     Compute propagation kernel for propagation with no Fresnel approximation.
diff --git a/src/chromatix/functional/pupils.py b/src/chromatix/functional/pupils.py
index b671285..8689ae5 100644
--- a/src/chromatix/functional/pupils.py
+++ b/src/chromatix/functional/pupils.py
@@ -1,10 +1,12 @@
+from chromatix.typing import ScalarLike
+
 from ..field import Field
 from ..utils import l2_sq_norm, linf_norm
 
 __all__ = ["circular_pupil", "square_pupil"]
 
 
-def circular_pupil(field: Field, w: float) -> Field:
+def circular_pupil(field: Field, w: ScalarLike) -> Field:
     """Applies circular pupil with diameter ``w`` to ``field``."""
     mask = (
         l2_sq_norm(field.grid)
@@ -13,7 +15,7 @@ def circular_pupil(field: Field, w: float) -> Field:
     return field * mask
 
 
-def square_pupil(field: Field, w: float) -> Field:
+def square_pupil(field: Field, w: ScalarLike) -> Field:
     """Applies square pupil with side length ``w`` to ``field``."""
     mask = linf_norm(field.grid) <= (field.spectrum[..., 0, 0] / field.spectrum) * w / 2
     return field * mask
diff --git a/src/chromatix/functional/samples.py b/src/chromatix/functional/samples.py
index d8e3b93..f6e6e91 100644
--- a/src/chromatix/functional/samples.py
+++ b/src/chromatix/functional/samples.py
@@ -1,11 +1,10 @@
-from typing import Optional, Tuple, Union
-
 import jax
 import jax.numpy as jnp
-from chex import assert_equal_shape, assert_rank
+from chex import PRNGKey, assert_equal_shape, assert_rank
 from jax import Array
 
 from chromatix.field import crop, pad
+from chromatix.typing import ArrayLike, ScalarLike
 from chromatix.utils.fft import fft, ifft
 
 from ..field import ScalarField, VectorField
@@ -20,7 +19,7 @@
 
 
 def jones_sample(
-    field: VectorField, absorption: Array, dn: Array, thickness: Union[float, Array]
+    field: VectorField, absorption: ArrayLike, dn: ArrayLike, thickness: ScalarLike
 ) -> VectorField:
     """
     Perturbs an incoming ``VectorField`` as if it went through a thin sample
@@ -61,7 +60,7 @@ def jones_sample(
 
 
 def thin_sample(
-    field: ScalarField, absorption: Array, dn: Array, thickness: Union[float, Array]
+    field: ScalarField, absorption: ArrayLike, dn: ArrayLike, thickness: ScalarLike
 ) -> ScalarField:
     """
     Perturbs an incoming ``ScalarField`` as if it went through a thin sample
@@ -100,14 +99,14 @@ def thin_sample(
 
 def multislice_thick_sample(
     field: ScalarField,
-    absorption_stack: Array,
-    dn_stack: Array,
-    n: float,
-    thickness_per_slice: float,
+    absorption_stack: ArrayLike,
+    dn_stack: ArrayLike,
+    n: ScalarLike,
+    thickness_per_slice: ScalarLike,
     N_pad: int,
-    propagator: Optional[Array] = None,
-    kykx: Union[Array, Tuple[float, float]] = (0.0, 0.0),
-    reverse_propagate_distance: Optional[float] = None,
+    propagator: ArrayLike | None = None,
+    kykx: ArrayLike | tuple[float, float] = (0.0, 0.0),
+    reverse_propagate_distance: ScalarLike | None = None,
 ) -> ScalarField:
     """
     Perturbs incoming ``ScalarField`` as if it went through a thick sample. The
@@ -166,17 +165,17 @@ def multislice_thick_sample(
 
 def fluorescent_multislice_thick_sample(
     field: ScalarField,
-    fluorescence_stack: Array,
-    dn_stack: Array,
-    n: float,
-    thickness_per_slice: float,
+    fluorescence_stack: ArrayLike,
+    dn_stack: ArrayLike,
+    n: ScalarLike,
+    thickness_per_slice: ScalarLike,
     N_pad: int,
-    key: jax.random.PRNGKey,
+    key: PRNGKey,
     num_samples: int = 1,
-    propagator_forward: Optional[Array] = None,
-    propagator_backward: Optional[Array] = None,
-    kykx: Union[Array, Tuple[float, float]] = (0.0, 0.0),
-) -> Array:
+    propagator_forward: ArrayLike | None = None,
+    propagator_backward: ArrayLike | None = None,
+    kykx: ArrayLike | tuple[float, float] = (0.0, 0.0),
+) -> ScalarField:
     """
     Perturbs incoming ``ScalarField`` as if it went through a thick,
     transparent, and fluorescent sample, i.e. a sample consisting of some
@@ -288,7 +287,7 @@ def _sample(i, field_and_intensity_stack):
 
 
 # depolarised wave
-def PTFT(k, km: Array) -> Array:
+def PTFT(k: ArrayLike, km: ScalarLike) -> Array:
     Q = jnp.zeros((3, 3, *k.shape[1:]))
 
     # Setting diagonal
@@ -318,7 +317,7 @@ def bmatvec(a, b):
 
 
 def thick_sample_vector(
-    field: VectorField, scatter_potential: Array, dz: float, n: float
+    field: VectorField, scatter_potential: ArrayLike, dz: ScalarLike, n: ScalarLike
 ) -> VectorField:
     def P_op(u: Array) -> Array:
         phase_factor = jnp.exp(1j * kz * dz)
diff --git a/src/chromatix/functional/sensors.py b/src/chromatix/functional/sensors.py
index 4f22926..1ba8e94 100644
--- a/src/chromatix/functional/sensors.py
+++ b/src/chromatix/functional/sensors.py
@@ -1,24 +1,26 @@
-from typing import Callable, Literal, Optional, Union
+from typing import Callable, Literal
 
 import jax.numpy as jnp
-from chex import Array, PRNGKey
-from jax import vmap
+from chex import PRNGKey
+from jax import Array, vmap
 from jax.lax import psum
 
-from ..field import Field
+from chromatix import ScalarField, VectorField
+from chromatix.typing import ScalarLike
+
 from ..ops import approximate_shot_noise, shot_noise
 
 __all__ = ["basic_sensor"]
 
 
 def basic_sensor(
-    sensor_input: Union[Field, Array],
-    shot_noise_mode: Optional[Literal["approximate", "poisson"]] = None,
-    resample_fn: Optional[Callable[[Array, float], Array]] = None,
-    reduce_axis: Optional[int] = None,
-    reduce_parallel_axis_name: Optional[str] = None,
-    input_spacing: Optional[Union[float, Array]] = None,
-    noise_key: Optional[PRNGKey] = None,
+    sensor_input: ScalarField | VectorField | Array,
+    shot_noise_mode: Literal["approximate", "poisson"] | None = None,
+    resample_fn: Callable[[Array, ScalarLike], Array] | None = None,
+    reduce_axis: int | None = None,
+    reduce_parallel_axis_name: str | None = None,
+    input_spacing: ScalarLike | None = None,
+    noise_key: PRNGKey | None = None,
 ) -> Array:
     """
     Produces an intensity image from an incoming ``Field`` or intensity
@@ -41,16 +43,17 @@ def basic_sensor(
             spacing of the input to be used for resampling by the sensor.
         noise_key: If provided, will be used to generate the shot noise.
     """
-    if isinstance(sensor_input, Field):
+    if isinstance(sensor_input, (ScalarField, VectorField)):
         intensity = sensor_input.intensity
         # WARNING(dd): @copypaste(Microscope) Assumes that field has same
         # spacing at all wavelengths when calculating intensity!
         input_spacing = sensor_input.dx[..., 0, 0].squeeze()
     else:
         intensity = sensor_input
+
     if resample_fn is not None:
         assert input_spacing is not None, "Must provide input_spacing for intensity"
-    if resample_fn is not None:
+    if (resample_fn is not None) and (input_spacing is not None):
         for i in range(sensor_input.ndim - 4):
             resample_fn = vmap(resample_fn, in_axes=(0, None))
         image = resample_fn(intensity, input_spacing)
diff --git a/src/chromatix/functional/sources.py b/src/chromatix/functional/sources.py
index 6ecede4..fbe7928 100644
--- a/src/chromatix/functional/sources.py
+++ b/src/chromatix/functional/sources.py
@@ -1,14 +1,18 @@
-from typing import Callable, Optional, Tuple, Union
-import numpy as np
+from typing import Callable
+
 import jax.numpy as jnp
-from chex import Array, assert_axis_dimension, assert_equal_shape
-from chromatix.field import Field, ScalarField, VectorField
+import numpy as np
+from chex import assert_axis_dimension, assert_equal_shape
+
+from chromatix import Field, ScalarField, VectorField
+from chromatix.typing import ArrayLike, ScalarLike
 from chromatix.utils import l2_sq_norm
 from chromatix.utils.shapes import (
     _broadcast_1d_to_grid,
     _broadcast_1d_to_innermost_batch,
     _broadcast_1d_to_polarization,
 )
+
 from .pupils import circular_pupil
 
 __all__ = [
@@ -19,19 +23,23 @@
 ]
 
 
+# We need this alias for typing to pass
+FieldPupil = Callable[[Field], Field]
+
+
 def point_source(
-    shape: Tuple[int, int],
-    dx: Union[float, Array],
-    spectrum: Union[float, Array],
-    spectral_density: Union[float, Array],
-    z: float,
-    n: float,
-    power: float = 1.0,
-    amplitude: Union[float, Array] = 1.0,
-    pupil: Optional[Callable[[ScalarField], ScalarField]] = None,
+    shape: tuple[int, int],
+    dx: ScalarLike,
+    spectrum: ScalarLike,
+    spectral_density: ScalarLike,
+    z: ScalarLike,
+    n: ScalarLike,
+    power: ScalarLike = 1.0,
+    amplitude: ScalarLike = 1.0,
+    pupil: FieldPupil | None = None,
     scalar: bool = True,
-    epsilon: float = np.finfo(np.float32).eps,
-) -> Field:
+    epsilon: float = float(np.finfo(np.float32).eps),
+) -> ScalarField | VectorField:
     """
     Generates field due to point source a distance ``z`` away.
 
@@ -76,19 +84,19 @@ def point_source(
 
 
 def objective_point_source(
-    shape: Tuple[int, int],
-    dx: Union[float, Array],
-    spectrum: Union[float, Array],
-    spectral_density: Union[float, Array],
-    z: float,
-    f: float,
-    n: float,
-    NA: float,
-    power: float = 1.0,
-    amplitude: Union[float, Array] = 1.0,
-    offset: Union[Array, Tuple[float, float]] = (0.0, 0.0),
+    shape: tuple[int, int],
+    dx: ScalarLike,
+    spectrum: ScalarLike,
+    spectral_density: ScalarLike,
+    z: ScalarLike,
+    f: ScalarLike,
+    n: ScalarLike,
+    NA: ScalarLike,
+    power: ScalarLike = 1.0,
+    amplitude: ScalarLike = 1.0,
+    offset: ArrayLike | tuple[float, float] = (0.0, 0.0),
     scalar: bool = True,
-) -> Field:
+) -> ScalarField | VectorField:
     """
     Generates field due to a point source defocused by an amount ``z`` away from
     the focal plane, just after passing through a lens with focal length ``f``
@@ -132,21 +140,21 @@ def objective_point_source(
     u = amplitude * -1j / L**2 * jnp.exp(1j * phase)
     field = field.replace(u=u)
     D = 2 * f * NA / n
-    field = circular_pupil(field, D)
+    field = circular_pupil(field, D)  # type: ignore
     return field * jnp.sqrt(power / field.power)
 
 
 def plane_wave(
-    shape: Tuple[int, int],
-    dx: Union[float, Array],
-    spectrum: Union[float, Array],
-    spectral_density: Union[float, Array],
-    power: float = 1.0,
-    amplitude: Union[float, Array] = 1.0,
-    kykx: Union[Array, Tuple[float, float]] = (0.0, 0.0),
-    pupil: Optional[Callable[[Field], Field]] = None,
+    shape: tuple[int, int],
+    dx: ScalarLike,
+    spectrum: ScalarLike,
+    spectral_density: ScalarLike,
+    power: ScalarLike = 1.0,
+    amplitude: ScalarLike = 1.0,
+    kykx: ArrayLike | tuple[float, float] = (0.0, 0.0),
+    pupil: FieldPupil | None = None,
     scalar: bool = True,
-) -> Field:
+) -> ScalarField | VectorField:
     """
     Generates plane wave of given ``power``.
 
@@ -190,15 +198,15 @@ def plane_wave(
 
 
 def generic_field(
-    dx: Union[float, Array],
-    spectrum: Union[float, Array],
-    spectral_density: Union[float, Array],
-    amplitude: Array,
-    phase: Array,
-    power: Optional[float] = 1.0,
-    pupil: Optional[Callable[[ScalarField], ScalarField]] = None,
+    dx: ScalarLike,
+    spectrum: ScalarLike,
+    spectral_density: ScalarLike,
+    amplitude: ArrayLike,
+    phase: ArrayLike,
+    power: ScalarLike = 1.0,
+    pupil: FieldPupil | None = None,
     scalar: bool = True,
-) -> Field:
+) -> ScalarField | VectorField:
     """
     Generates field with arbitrary ``phase`` and ``amplitude``.
 
@@ -228,7 +236,7 @@ def generic_field(
     assert_axis_dimension(amplitude, -1, vectorial_dimension)
     assert_axis_dimension(phase, -1, vectorial_dimension)
     assert_equal_shape([amplitude, phase])
-    u = amplitude * jnp.exp(1j * phase)
+    u = jnp.array(amplitude) * jnp.exp(1j * phase)
     field = create(dx, spectrum, spectral_density, u=u)
     if pupil is not None:
         field = pupil(field)
diff --git a/src/chromatix/ops/filters.py b/src/chromatix/ops/filters.py
index 619c693..6f892bb 100644
--- a/src/chromatix/ops/filters.py
+++ b/src/chromatix/ops/filters.py
@@ -1,7 +1,8 @@
-from typing import Optional, Sequence, Tuple
+from typing import Sequence
 
-from chex import Array
+from jax import Array
 
+from chromatix.typing import ArrayLike
 from chromatix.utils import _broadcast_2d_to_spatial, gaussian_kernel
 
 from .ops import fourier_convolution
@@ -10,10 +11,10 @@
 
 
 def high_pass_filter(
-    data: Array,
+    data: ArrayLike,
     sigma: Sequence[float],
-    axes: Tuple[int, int] = (1, 2),
-    kernel_shape: Optional[Sequence[int]] = None,
+    axes: tuple[int, int] = (1, 2),
+    kernel_shape: Sequence[int] | None = None,
 ) -> Array:
     """
     Performs a high pass filter on ``data``.
@@ -47,8 +48,8 @@ def high_pass_filter(
 def gaussian_filter(
     data: Array,
     sigma: Sequence[float],
-    axes: Tuple[int, int] = (1, 2),
-    kernel_shape: Optional[Sequence[int]] = None,
+    axes: tuple[int, int] = (1, 2),
+    kernel_shape: Sequence[int] | None = None,
 ) -> Array:
     """
     Performs a Gaussian filter on ``data``.
diff --git a/src/chromatix/ops/noise.py b/src/chromatix/ops/noise.py
index ee8ec13..2a16e8b 100644
--- a/src/chromatix/ops/noise.py
+++ b/src/chromatix/ops/noise.py
@@ -1,12 +1,14 @@
 from typing import Tuple
 
 import jax.numpy as jnp
-from chex import Array, PRNGKey
-from jax import custom_jvp, random
+from chex import PRNGKey
+from jax import Array, custom_jvp, random
+
+from chromatix.typing import ArrayLike
 
 
 @custom_jvp
-def approximate_shot_noise(key: PRNGKey, image: Array) -> Array:
+def approximate_shot_noise(key: PRNGKey, image: ArrayLike) -> Array:
     """
     Approximates Poisson shot noise using a Gaussian for differentiability.
     """
@@ -15,7 +17,7 @@ def approximate_shot_noise(key: PRNGKey, image: Array) -> Array:
 
 
 @approximate_shot_noise.defjvp
-def approximate_shotnoise_jvp(primals: Tuple, tangents: Tuple) -> Tuple:
+def approximate_shotnoise_jvp(primals: tuple, tangents: tuple) -> tuple:
     """
     Custom gradient for ``approximate_shot_noise``.
 
@@ -35,7 +37,7 @@ def approximate_shotnoise_jvp(primals: Tuple, tangents: Tuple) -> Tuple:
 
 
 @custom_jvp
-def shot_noise(key: PRNGKey, image: Array) -> Array:
+def shot_noise(key: PRNGKey, image: ArrayLike) -> Array:
     """
     Simulates Poisson shot noise whose gradient is approximated using
     the gradient of a Gaussian, just as if the simulation had been
diff --git a/src/chromatix/ops/ops.py b/src/chromatix/ops/ops.py
index c30aa67..b31c1bf 100644
--- a/src/chromatix/ops/ops.py
+++ b/src/chromatix/ops/ops.py
@@ -1,14 +1,18 @@
 from functools import partial
 
 import jax.numpy as jnp
-from chex import Array
-from jax import lax
+from jax import Array, lax
 
+from chromatix.typing import ArrayLike
 from chromatix.utils import next_order
 
 
 def fourier_convolution(
-    image: Array, kernel: Array, *, axes=(0, 1), fast_fft_shape: bool = True
+    image: ArrayLike,
+    kernel: ArrayLike,
+    *,
+    axes: tuple[int, int] = (0, 1),
+    fast_fft_shape: bool = True,
 ) -> Array:
     """
     Fourier convolution in 2D over the specified axes of an ``Array``.
diff --git a/src/chromatix/ops/quantization.py b/src/chromatix/ops/quantization.py
index 069d51b..2993d5d 100644
--- a/src/chromatix/ops/quantization.py
+++ b/src/chromatix/ops/quantization.py
@@ -1,8 +1,6 @@
-from typing import Optional, Tuple
-
 import jax
 import jax.numpy as jnp
-from chex import Array
+from jax import Array
 
 __all__ = ["binarize", "binarize_jvp", "quantize", "quantize_jvp"]
 
@@ -30,7 +28,7 @@ def binarize(x: Array, threshold: float = 0.5) -> Array:
 
 
 @binarize.defjvp
-def binarize_jvp(primals: Tuple, tangents: Tuple) -> Tuple:
+def binarize_jvp(primals: tuple, tangents: tuple) -> tuple:
     """
     Custom gradient for ``binarize``.
 
@@ -43,9 +41,7 @@ def binarize_jvp(primals: Tuple, tangents: Tuple) -> Tuple:
 
 
 @jax.custom_jvp
-def quantize(
-    x: Array, bit_depth: float, range: Optional[Tuple[int, int]] = None
-) -> Array:
+def quantize(x: Array, bit_depth: float, range: tuple[int, int] | None = None) -> Array:
     """
     Quantize the input ``x`` to the specified ``bit_depth``. Surrogate gradient
     approach [1] is used to adjust the bit depth differentiably.
@@ -74,7 +70,7 @@ def quantize(
 
 
 @quantize.defjvp
-def quantize_jvp(primals: Tuple, tangents: Tuple) -> Tuple:
+def quantize_jvp(primals: tuple, tangents: tuple) -> tuple:
     """
     Custom gradient for ``quantize``.
 
diff --git a/src/chromatix/ops/resample.py b/src/chromatix/ops/resample.py
index eb520ef..fbc97fd 100644
--- a/src/chromatix/ops/resample.py
+++ b/src/chromatix/ops/resample.py
@@ -1,13 +1,15 @@
-from typing import Callable, Tuple, Union
+from typing import Callable
 
 import jax.numpy as jnp
-from chex import Array
 from einops import reduce
+from jax import Array
 from jax.image import scale_and_translate
 
+from chromatix.typing import ScalarLike
+
 
 def pooling_downsample(
-    data: Array, window_size: Tuple[int, int], reduction: str = "mean"
+    data: Array, window_size: tuple[int, int], reduction: str = "mean"
 ) -> Array:
     """
     Wrapper for downsampling input of shape `(B... H W C P)` along `(H W)`.
@@ -34,10 +36,10 @@ def pooling_downsample(
 
 
 def init_plane_resample(
-    out_shape: Tuple[int, ...],
-    out_spacing: Union[float, Array],
+    out_shape: tuple[int, ...],
+    out_spacing: ScalarLike,
     resampling_method: str = "linear",
-) -> Callable[[Array, float], Array]:
+) -> Callable[[Array, ScalarLike], Array]:
     """
     Returns a function that resamples 2D planes to the specified output shape
     and spacing.
@@ -64,7 +66,7 @@ def init_plane_resample(
     ), "Spacing is either a float or array of shape (2,) for non-square pixels"
     if resampling_method == "pool":
 
-        def op(x: Array, in_spacing: Union[float, Array]) -> Array:
+        def op(x: Array, in_spacing: ScalarLike) -> Array:
             return reduce(
                 x,
                 "(h hf) (w wf) ... -> h w ...",
@@ -75,7 +77,7 @@ def op(x: Array, in_spacing: Union[float, Array]) -> Array:
 
     else:
 
-        def op(x: Array, in_spacing: Union[float, Array]) -> Array:
+        def op(x: Array, in_spacing: ScalarLike) -> Array:
             in_spacing = jnp.atleast_1d(in_spacing).squeeze()
             assert (
                 in_spacing.size <= 2
diff --git a/src/chromatix/systems/microscopes.py b/src/chromatix/systems/microscopes.py
index efd28b1..51c5759 100644
--- a/src/chromatix/systems/microscopes.py
+++ b/src/chromatix/systems/microscopes.py
@@ -1,15 +1,17 @@
 from __future__ import annotations
 
-from typing import Any, Callable, Tuple, Union
+from typing import Any, Callable
 
 import jax.numpy as jnp
-from chex import Array, PRNGKey
+from chex import PRNGKey
 from flax import linen as nn
+from jax import Array
 
+from chromatix.typing import ArrayLike, ScalarLike
 from chromatix.utils import sigmoid_taper
 
 from ..elements import FFLens, ObjectivePointSource, PhaseMask
-from ..field import Field
+from ..field import ScalarField, VectorField
 from ..ops import fourier_convolution
 from ..utils import center_crop
 from .optical_system import OpticalSystem
@@ -71,18 +73,18 @@ class Microscope(nn.Module):
             this to `False`. Defaults to `True`.
     """
 
-    system_psf: Callable[[Microscope], Union[Field, Array]]
+    system_psf: Callable[[Microscope], ScalarField | VectorField | Array]
     sensor: nn.Module
-    f: float
-    n: float
-    NA: float
-    spectrum: Array
-    spectral_density: Array
+    f: ScalarLike
+    n: ScalarLike
+    NA: ScalarLike
+    spectrum: ScalarLike
+    spectral_density: ScalarLike
     padding_ratio: float = 0
     taper_width: float = 0
     fast_fft_shape: bool = True
 
-    def __call__(self, sample: Array, *args: Any, **kwargs: Any) -> Array:
+    def __call__(self, sample: ArrayLike, *args: Any, **kwargs: Any) -> Array:
         """
         Computes PSF and convolves PSF with ``data`` to simulate imaging.
 
@@ -101,12 +103,15 @@ def __call__(self, sample: Array, *args: Any, **kwargs: Any) -> Array:
         psf = self._process_psf(system_psf, ndim, spatial_dims)
         return self.image(sample, psf, axes=spatial_dims)
 
-    def psf(self, *args: Any, **kwargs: Any) -> Union[Field, Array]:
+    def psf(self, *args: Any, **kwargs: Any) -> ScalarField | VectorField | Array:
         """Computes PSF of system, taking any necessary arguments."""
         return self.system_psf(self, *args, **kwargs)
 
     def _process_psf(
-        self, system_psf: Union[Field, Array], ndim: int, spatial_dims: Tuple[int, int]
+        self,
+        system_psf: ScalarField | VectorField | Array,
+        ndim: int,
+        spatial_dims: tuple[int, int],
     ) -> Array:
         """
         Prepare PSF to be convolved with a sample by doing the following:
@@ -125,7 +130,7 @@ def _process_psf(
             padding = (0, 0)
         # WARNING(dd): Assumes that field has same spacing at all wavelengths
         # when calculating intensity!
-        if isinstance(system_psf, Field):
+        if isinstance(system_psf, (ScalarField | VectorField)):
             psf = system_psf.intensity
             spacing = system_psf.dx[..., 0, 0].squeeze()
         else:
@@ -137,16 +142,18 @@ def _process_psf(
                 ]
             )
         if self.padding_ratio > 0:
-            pad_spec = [None for _ in range(ndim)]
+            pad_spec: list[Any] = [None for _ in range(ndim)]
             pad_spec[spatial_dims[0]] = padding[0] // 2
             pad_spec[spatial_dims[1]] = padding[1] // 2
             psf = center_crop(psf, pad_spec)
         if self.taper_width > 0:
-            psf = psf * sigmoid_taper(unpadded_shape, self.taper_width, ndim=ndim)
+            psf = psf * sigmoid_taper(unpadded_shape, self.taper_width, ndim=ndim)  # type: ignore
         psf = self.sensor.resample(psf, spacing)
         return psf
 
-    def image(self, sample: Array, psf: Array, axes: Tuple[int, int] = (1, 2)) -> Array:
+    def image(
+        self, sample: ArrayLike, psf: ArrayLike, axes: tuple[int, int] = (1, 2)
+    ) -> Array:
         """
         Computes image or batch of images using the specified ``psf`` and
         ``sample``. Assumes that both the ``sample`` and ``psf`` have already
@@ -156,7 +163,9 @@ def image(self, sample: Array, psf: Array, axes: Tuple[int, int] = (1, 2)) -> Ar
             sample: The sample volume to image with of shape `(B... H W 1 1)`.
             psf: The PSF intensity volume to image with of shape `(B... H W 1 1)`.
         """
-        image = fourier_convolution(sample, psf, axes=axes, fast_fft_shape=self.fast_fft_shape)
+        image = fourier_convolution(
+            sample, psf, axes=axes, fast_fft_shape=self.fast_fft_shape
+        )
         # NOTE(dd): By this point, the image should already be at the same
         # spacing as the sensor. Any resampling to the pixels of the sensor
         # should already have happened to the PSF.
@@ -183,12 +192,12 @@ class Optical4FSystemPSF(nn.Module):
             initialization function (e.g. using trainable).
     """
 
-    shape: Tuple[int, int]
-    spacing: float
-    phase: Union[Array, Callable[[PRNGKey, Tuple[int, ...]], Array]]
+    shape: tuple[int, int]
+    spacing: ScalarLike
+    phase: ArrayLike | Callable[[PRNGKey, tuple[int, int], Array, Array], Array]
 
     @nn.compact
-    def __call__(self, microscope: Microscope, z: Array) -> Field:
+    def __call__(self, microscope: Microscope, z: ArrayLike) -> Array:
         padding = tuple(int(s * microscope.padding_ratio) for s in self.shape)
         padded_shape = tuple(s + p for s, p in zip(self.shape, padding))
         required_spacing = self.compute_required_spacing(
@@ -201,7 +210,7 @@ def __call__(self, microscope: Microscope, z: Array) -> Field:
         system = OpticalSystem(
             [
                 ObjectivePointSource(
-                    padded_shape,
+                    padded_shape,  # type: ignore
                     required_spacing,
                     microscope.spectrum,
                     microscope.spectral_density,
@@ -218,6 +227,10 @@ def __call__(self, microscope: Microscope, z: Array) -> Field:
 
     @staticmethod
     def compute_required_spacing(
-        height: int, output_spacing: float, f: float, n: float, spectrum: Array
-    ) -> float:
+        height: int,
+        output_spacing: ScalarLike,
+        f: ScalarLike,
+        n: ScalarLike,
+        spectrum: ScalarLike,
+    ) -> ScalarLike:
         return f * spectrum / (n * height * output_spacing)
diff --git a/src/chromatix/systems/optical_system.py b/src/chromatix/systems/optical_system.py
index ccdea4a..eb6529e 100644
--- a/src/chromatix/systems/optical_system.py
+++ b/src/chromatix/systems/optical_system.py
@@ -1,9 +1,9 @@
 from __future__ import annotations
 
-from typing import Any, Callable, Sequence, Union
+from typing import Any, Callable, Sequence
 
-from chex import Array
 from flax import linen as nn
+from jax import Array
 
 from ..field import Field
 
@@ -29,7 +29,7 @@ class OpticalSystem(nn.Module):
     elements: Sequence[Callable]
 
     @nn.compact
-    def __call__(self, *args: Any, **kwargs: Any) -> Union[Field, Array]:
+    def __call__(self, *args: Any, **kwargs: Any) -> Field | Array:
         """Returns the result of calling all elements in sequence."""
         field = self.elements[0](*args, **kwargs)  # allow field to be initialized
         for element in self.elements[1:]:
diff --git a/src/chromatix/typing.py b/src/chromatix/typing.py
new file mode 100644
index 0000000..2fcad92
--- /dev/null
+++ b/src/chromatix/typing.py
@@ -0,0 +1,6 @@
+from jax import Array
+from numpy import ndarray
+
+ArrayLike = Array | ndarray
+ScalarLike = ArrayLike | float | int
+ComplexScalarLike = ArrayLike | float | int | complex
diff --git a/src/chromatix/utils/__init__.py b/src/chromatix/utils/__init__.py
index f9a1e12..d23d65b 100644
--- a/src/chromatix/utils/__init__.py
+++ b/src/chromatix/utils/__init__.py
@@ -1,4 +1,4 @@
-from .data import *
+from ..data.data import *
 from .initializers import *
 from .shapes import *
 from .utils import *
diff --git a/src/chromatix/utils/fft.py b/src/chromatix/utils/fft.py
index 9ac0692..6c2e9b0 100644
--- a/src/chromatix/utils/fft.py
+++ b/src/chromatix/utils/fft.py
@@ -1,11 +1,12 @@
 from functools import partial
-from typing import Tuple
 
 import jax.numpy as jnp
-from chex import Array
+from jax import Array
 
+from chromatix.typing import ArrayLike
 
-def fft(x: Array, axes: Tuple[int, int] = (1, 2), shift: bool = False) -> Array:
+
+def fft(x: ArrayLike, axes: tuple[int, int] = (1, 2), shift: bool = False) -> Array:
     """
     Computes ``fft2`` for input of shape `(B... H W C)`.
     If shift is true, first applies ``ifftshift``, than an ``fftshift`` to
@@ -20,7 +21,7 @@ def fft(x: Array, axes: Tuple[int, int] = (1, 2), shift: bool = False) -> Array:
         return fft(x)
 
 
-def ifft(x: Array, axes: Tuple[int, int] = (1, 2), shift: bool = False) -> Array:
+def ifft(x: ArrayLike, axes: tuple[int, int] = (1, 2), shift: bool = False) -> Array:
     """
     Computes ``ifft2`` for input of shape `(B... H W C)`.
     If shift is true, first applies ``ifftshift``, than an ``fftshift`` to
diff --git a/src/chromatix/utils/initializers.py b/src/chromatix/utils/initializers.py
index 9f117db..69cf7f2 100644
--- a/src/chromatix/utils/initializers.py
+++ b/src/chromatix/utils/initializers.py
@@ -1,11 +1,13 @@
 import math
-from typing import Sequence, Tuple
+from typing import Sequence
 
 import jax.numpy as jnp
-from chex import Array
 from einops import rearrange
+from jax import Array
 from scipy.special import comb  # type: ignore
 
+from chromatix.typing import ScalarLike
+
 from .utils import create_grid, grid_spatial_to_pupil
 
 __all__ = [
@@ -17,7 +19,7 @@
 ]
 
 
-def flat_phase(shape: Tuple[int, int], *args, value: float = 0.0) -> Array:
+def flat_phase(shape: tuple[int, int], *args, value: ScalarLike = 0.0) -> Array:
     """
     Computes a flat mask (one with constant value).
 
@@ -30,14 +32,14 @@ def flat_phase(shape: Tuple[int, int], *args, value: float = 0.0) -> Array:
 
 
 def potato_chip(
-    shape: Tuple[int, ...],
-    spacing: float,
-    wavelength: float,
-    n: float,
-    f: float,
-    NA: float,
-    d: float = 50.0,
-    C0: float = -146.7,
+    shape: tuple[int, int],
+    spacing: ScalarLike,
+    wavelength: ScalarLike,
+    n: ScalarLike,
+    f: ScalarLike,
+    NA: ScalarLike,
+    d: ScalarLike = 50.0,
+    C0: ScalarLike = -146.7,
 ) -> Array:
     """
     Computes the "potato chip" phase mask described by [1].
@@ -77,15 +79,15 @@ def potato_chip(
 
 
 def seidel_aberrations(
-    shape: Tuple[int, int],
-    spacing: float,
-    wavelength: float,
-    n: float,
-    f: float,
-    NA: float,
+    shape: tuple[int, int],
+    spacing: ScalarLike,
+    wavelength: ScalarLike,
+    n: ScalarLike,
+    f: ScalarLike,
+    NA: ScalarLike,
     coefficients: Sequence[float],
-    u: float = 0,
-    v: float = 0,
+    u: ScalarLike = 0,
+    v: ScalarLike = 0,
 ) -> Array:
     """
     Computes the Seidel phase polynomial described by [1]. Accounts for spatially
@@ -137,12 +139,12 @@ def seidel_aberrations(
 
 
 def zernike_aberrations(
-    shape: Tuple[int, int],
-    spacing: float,
-    wavelength: float,
-    n: float,
-    f: float,
-    NA: float,
+    shape: tuple[int, int],
+    spacing: ScalarLike,
+    wavelength: ScalarLike,
+    n: ScalarLike,
+    f: ScalarLike,
+    NA: ScalarLike,
     ansi_indices: Sequence[int],
     coefficients: Sequence[float],
     normalization: bool = True,
@@ -234,12 +236,12 @@ def R(rho):
 
 
 def defocused_ramps(
-    shape: Tuple[int, int],
-    spacing: float,
-    wavelength: float,
-    n: float,
-    f: float,
-    NA: float,
+    shape: tuple[int, int],
+    spacing: ScalarLike,
+    wavelength: ScalarLike,
+    n: ScalarLike,
+    f: ScalarLike,
+    NA: ScalarLike,
     num_ramps: int = 6,
     delta: Sequence[float] = [2374.0] * 6,
     defocus: Sequence[float] = [-50.0, 150.0, -100.0, 50.0, -150.0, 100.0],
@@ -308,7 +310,7 @@ def ramp(center, theta_bounds, delta_ramp, ramp_defocus):
             grid[0] - jnp.sin(center) * defocus_center
         ) ** 2
         phase += ramp_mask * (ramp_defocus * ramp_quadratic)
-        phase -= ramp_mask * jnp.where(ramp_mask > 0, phase, 0).mean()
+        phase -= ramp_mask * jnp.where(ramp_mask > 0, phase, 0).mean()  # type: ignore
         return phase
 
     for ramp_idx in range(num_ramps):
diff --git a/src/chromatix/utils/shapes.py b/src/chromatix/utils/shapes.py
index b45cd5b..bc5f676 100644
--- a/src/chromatix/utils/shapes.py
+++ b/src/chromatix/utils/shapes.py
@@ -1,8 +1,8 @@
-from typing import Union
-
 import jax.numpy as jnp
-from chex import Array
 from einops import rearrange
+from jax import Array
+
+from chromatix.typing import ArrayLike, ScalarLike
 
 __all__ = [
     "_broadcast_1d_to_channels",
@@ -15,34 +15,34 @@
 ]
 
 
-def _broadcast_1d_to_channels(x: Union[float, Array], ndim: int) -> Array:
+def _broadcast_1d_to_channels(x: ScalarLike, ndim: int) -> Array:
     """Broadcast 1D array of size `C` to `(B... H W C [1 | 3])`."""
     shape_spec = "c -> " + ("1 " * (ndim - 2)) + "c 1"
     return rearrange(jnp.atleast_1d(x), shape_spec)
 
 
-def _broadcast_1d_to_polarization(x: Union[float, Array], ndim: int) -> Array:
+def _broadcast_1d_to_polarization(x: ScalarLike, ndim: int) -> Array:
     """Broadcast 1D array of size `P` to `(B... H W C [1 | 3])`."""
     shape_spec = "p -> " + ("1 " * (ndim - 1)) + "p"
     return rearrange(jnp.atleast_1d(x), shape_spec)
 
 
-def _broadcast_1d_to_innermost_batch(x: Union[float, Array], ndim: int) -> Array:
+def _broadcast_1d_to_innermost_batch(x: ScalarLike, ndim: int) -> Array:
     """Broadcast 1D array of size `B` to left of `(H W)` in `(B... H W C [1 | 3])`."""
     shape_spec = "b ->" + " 1" * (ndim - 5) + " b 1 1 1 1"
     return rearrange(jnp.atleast_1d(x), shape_spec)
 
 
-def _broadcast_1d_to_grid(x: Union[float, Array], ndim: int) -> Array:
+def _broadcast_1d_to_grid(x: ArrayLike | tuple[float, float], ndim: int) -> Array:
     """
     Broadcast 1D array of size `2` to `(2 B... H W C 1)`.
     Useful for vectorial ops on grids.
     """
     shape_spec = "d ->" + "d" + " 1" * (ndim - 4) + " 1 1 1 1"
-    return rearrange(jnp.atleast_1d(x), shape_spec, d=2)
+    return rearrange(jnp.atleast_1d(jnp.array(x)), shape_spec, d=2)
 
 
-def _broadcast_2d_to_grid(x: Union[float, Array], ndim: int) -> Array:
+def _broadcast_2d_to_grid(x: Array, ndim: int) -> Array:
     """
     Broadcast 2D array of shape `2 C` to `(2 B... H W C [1 | 3])`.
     Useful for vectorial ops on grids.
@@ -51,7 +51,7 @@ def _broadcast_2d_to_grid(x: Union[float, Array], ndim: int) -> Array:
     return rearrange(x, shape_spec, d=2)
 
 
-def _squeeze_grid_to_2d(x: Union[float, Array], ndim: int) -> Array:
+def _squeeze_grid_to_2d(x: Array, ndim: int) -> Array:
     """
     Squeeze array of shape `(2 B... H W C [1 | 3])` to 2D array of shape `2 C`.
     Useful for vectorial ops on grids.
diff --git a/src/chromatix/utils/utils.py b/src/chromatix/utils/utils.py
index 53e04c5..f50b854 100644
--- a/src/chromatix/utils/utils.py
+++ b/src/chromatix/utils/utils.py
@@ -1,11 +1,13 @@
-from typing import Optional, Sequence, Tuple, Union
+from typing import Sequence
 
 import flax.linen as nn
 import jax.numpy as jnp
 import numpy as np
-from chex import Array
 from einops import rearrange
-from scipy.ndimage import distance_transform_edt  # type: ignore
+from jax import Array
+from scipy.ndimage import distance_transform_edt
+
+from chromatix.typing import ArrayLike, ScalarLike
 
 from .shapes import _broadcast_2d_to_spatial
 
@@ -14,7 +16,7 @@ def next_order(val: int) -> int:
     return int(2 ** np.ceil(np.log2(val)))
 
 
-def center_pad(u: jnp.ndarray, pad_width: Sequence[int], cval: float = 0) -> Array:
+def center_pad(u: ArrayLike, pad_width: Sequence[int], cval: float = 0) -> Array:
     """
     Symmetrically pads ``u`` with lengths specified per axis in ``n_padding``,
     which should be iterable and have the same size as ``u.ndims``.
@@ -23,7 +25,7 @@ def center_pad(u: jnp.ndarray, pad_width: Sequence[int], cval: float = 0) -> Arr
     return jnp.pad(u, pad, constant_values=cval)
 
 
-def center_crop(u: jnp.ndarray, crop_length: Sequence[int]) -> Array:
+def center_crop(u: Array, crop_length: Sequence[int]) -> Array:
     """
     Symmetrically crops ``u`` with lengths specified per axis in
     ``crop_length``, which should be iterable with same size as ``u.ndims``.
@@ -34,7 +36,7 @@ def center_crop(u: jnp.ndarray, crop_length: Sequence[int]) -> Array:
 
 
 def gaussian_kernel(
-    sigma: Sequence[float], truncate: float = 4.0, shape: Optional[Sequence[int]] = None
+    sigma: Sequence[float], truncate: float = 4.0, shape: Sequence[int] | None = None
 ) -> Array:
     """
     Creates ND Gaussian kernel of given ``sigma``.
@@ -74,13 +76,13 @@ def gaussian_kernel(
     return phi / phi.sum()
 
 
-def sigmoid_taper(shape: Tuple[int, int], width: float, ndim: int = 5) -> Array:
+def sigmoid_taper(shape: tuple[int, int], width: float, ndim: int = 5) -> Array:
     dist = distance_transform_edt(np.pad(np.ones((shape[0] - 2, shape[1] - 2)), 1))
-    taper = 2 * (nn.sigmoid(dist / width) - 0.5)
+    taper = 2 * (nn.sigmoid(dist / width) - 0.5)  # type: ignore - it's an array!
     return _broadcast_2d_to_spatial(taper, ndim)
 
 
-def create_grid(shape: Tuple[int, int], spacing: Union[float, Array]) -> Array:
+def create_grid(shape: tuple[int, int], spacing: ScalarLike) -> Array:
     """
     Args:
         shape: The shape of the grid, described as a tuple of
@@ -106,27 +108,29 @@ def create_grid(shape: Tuple[int, int], spacing: Union[float, Array]) -> Array:
     return grid
 
 
-def grid_spatial_to_pupil(grid: Array, f: float, NA: float, n: float) -> Array:
+def grid_spatial_to_pupil(
+    grid: Array, f: ScalarLike, NA: ScalarLike, n: ScalarLike
+) -> Array:
     R = f * NA / n  # pupil radius
     return grid / R
 
 
-def l2_sq_norm(a: Array, axis: Union[int, Tuple[int, ...]] = 0) -> Array:
+def l2_sq_norm(a: Array, axis: int | tuple[int, ...] = 0) -> Array:
     """Sum of squares, i.e. `x**2 + y**2`."""
     return jnp.sum(a**2, axis=axis)
 
 
-def l2_norm(a: Array, axis: Union[int, Tuple[int, ...]] = 0) -> Array:
+def l2_norm(a: Array, axis: int | tuple[int, ...] = 0) -> Array:
     """Square root of ``l2_sq_norm``, i.e. `sqrt(x**2 + y**2)`."""
     return jnp.sqrt(jnp.sum(a**2, axis=axis))
 
 
-def l1_norm(a: Array, axis: Union[int, Tuple[int, ...]] = 0) -> Array:
+def l1_norm(a: Array, axis: int | tuple[int, ...] = 0) -> Array:
     """Sum absolute value, i.e. `|x| + |y|`."""
     return jnp.sum(jnp.abs(a), axis=axis)
 
 
-def linf_norm(a: Array, axis: Union[int, Tuple[int, ...]] = 0) -> Array:
+def linf_norm(a: Array, axis: int | tuple[int, ...] = 0) -> Array:
     """Max absolute value, i.e. `max(|x|, |y|)`."""
     return jnp.max(jnp.abs(a), axis=axis)
 


Test patch:
diff --git a/tests/test_sources.py b/tests/test_sources.py
index 8dda0da..f41c576 100644
--- a/tests/test_sources.py
+++ b/tests/test_sources.py
@@ -6,6 +6,7 @@
 import jax.numpy as jnp
 import pytest
 from chex import assert_shape
+from jax import Array
 
 
 class TestPlaneWave: