Type Checking

Type Checking#

TL;DR#

You can annotate the dtype and shape of a Quantity — checked statically, and at runtime. To also constrain the physical dimension of an argument, use the dimension-parametrized ParametricQuantity from the separate unxts.parametric package; see Dimension annotations for type checking.

In the following example we define a function over two 1-D, equally-shaped float arrays and return their elementwise ratio:

from jaxtyping import Float

import unxt as u

def velocity(
    x: Float[u.Quantity, "N"],
    t: Float[u.Quantity, "N"],
) -> Float[u.Quantity, "N"]:
    return x / t

For information on typing in Python see the built-in typing module. Refer to the jaxtyping library for how to annotate the dtype and shape of a Quantity, for example integer arrays or variable / context-dependent shapes. jaxtyping also powers unxt’s runtime typechecking, discussed next.

Runtime Type Checking#

Using jaxtyping,unxt supports runtime type checking, where type annotations are enforced during execution. This is very useful for finding and preventing type-related errors, like passing the wrong type of argument to a function or returning the wrong type of value. To enable runtime type checking on all of unxt, set the environment variable UNXT_ENABLE_RUNTIME_TYPECHECKING to beartype.beartype or any other runtime typecheck backend supported by jaxtyping.

# Enable runtime type checking
export UNXT_ENABLE_RUNTIME_TYPECHECKING="beartype.beartype"

Attention

We recommend enabling runtime type checking during development.
For normal use, try enabling and disabling runtime type checking to assess any performance impact.

The performance overhead associated with runtime type checking should be small but isn’t always – in particular it can affect the time for JAX to jit code. To turn off runtime type checking set the environment variable to None.

# Disable runtime type checking
export UNXT_ENABLE_RUNTIME_TYPECHECKING="None"

Absent the environment variable, this is the default.

Tip

You can set environment variables directly in Python. Execute the following before importing unxt (or any library that imports unxt).

import os

os.environ["UNXT_ENABLE_RUNTIME_TYPECHECKING"] = "beartype.beartype"

In the background unxt checks for the UNXT_ENABLE_RUNTIME_TYPECHECKING environment variable and passes it to jaxtyping’s import hook. jaxtyping also offers function-specific checking through the jaxtyped decorator.

Here’s an example:

>>> from jaxtyping import Shaped, jaxtyped
>>> from beartype import beartype as typechecker  # or use any supported typechecker

>>> import unxt as u

>>> @jaxtyped(typechecker=typechecker)
... def velocity(
...     x: Shaped[u.Quantity, "N"],
...     t: Shaped[u.Quantity, "N"],
... ) -> Shaped[u.Quantity, "N"]:
...     return x / t

>>> x = u.Q([2.], "m")
>>> t = u.Q([1.], "s")

>>> velocity(x, t)
Quantity(Array([2.], dtype=float32), unit='m / s')

Dimension annotations#

The default unxt.quantity.Quantity carries dtype and shape in its type, but no dimension — so Quantity[<dimension>] does nothing (it is informational only). To additionally check the physical dimension of an argument (e.g. up.PQ["length"]), use ParametricQuantity from the separate unxts.parametric package, which encodes the dimension in its type. For construction, dimension inference/checking, and the parametric-class theory, see Dimension annotations for type checking.