unxt.experimental 🧪

unxt.experimental 🧪#

Experimental features.

Warning

These features may be removed or changed in the future without notice.

On some occasions JAX’s automatic differentiation functions do not work well with quantities. This is checked by enabling runtime type-checking (see the docs), which will raise an error if a quantity’s units do not match the expected input / output units of a function. In these cases, you can use the functions in this module to provide the units to the automatic differentiation functions. Instead of directly propagating the units through the automatic differentiation functions, the units are stripped and re-applied, while also being provided within the function being AD’d.

To import this experimental module

>>> from unxt import experimental
unxt._src.experimental.grad(fun: Callable[[Unpack[Args]], R], argnums: int = 0, *, units: tuple[UnitBase | FunctionUnitBase | str | None, ...])#

Gradient of a function with units.

In general, if you can use quax.quaxify(jax.grad(func)) (or the syntactic sugar quaxed.grad(func)), that’s the better option! The difference from those functions is how units are handled. quaxify will directly propagate the units through the automatic differentiation functions. But sometimes that doesn’t work and we need to strip the units and re-apply them. This function does that, using the units kwarg.

See also

jax.grad

The original JAX gradient function.

Examples

>>> import jax.numpy as jnp
>>> import unxt as u

>>> def cube_volume(x: u.Q["length"]) -> u.Q["volume"]:
...     return x**3

>>> grad_cube_volume = u.experimental.grad(cube_volume, units=("m",))
>>> grad_cube_volume(u.Q(2.0, "m"))
Quantity(Array(12., dtype=float32...), unit='m2')

Inputs are converted to units first, so a convertible input gives the same result (200 cm is 2 m):

>>> grad_cube_volume(u.Q(200.0, "cm"))
Quantity(Array(12., dtype=float32...), unit='m2')

A None entry in units marks an argument as a plain (unitless) value rather than a Quantity, so functions can mix the two:

>>> def scaled_area(distance, factor):
...     return distance**2 * factor

>>> g = u.experimental.grad(scaled_area, argnums=0, units=("m", None))
>>> g(u.Q(3.0, "m"), 2.0)
Quantity(Array(12., dtype=float32...), unit='m')
Parameters:
Return type:

Callable[[Unpack[TypeVarTuple]], TypeVar(R, bound= AbstractQuantity)]

unxt._src.experimental.hessian(fun: Callable[[Unpack[Args]], R], argnums: int = 0, *, units: tuple[UnitBase | FunctionUnitBase | str | None, ...])#

Hessian.

In general, if you can use quax.quaxify(jax.hessian(func)) (or the syntactic sugar quax.hessian(func)), that’s the better option! The difference from those functions is how units are handled. quaxify will directly propagate the units through the automatic differentiation functions. But sometimes that doesn’t work and we need to strip the units and re-apply them. This function does that, using the units kwarg.

See also

jax.hessian

The original JAX hessian function.

Examples

>>> import jax.numpy as jnp
>>> import unxt as u

>>> def cubbe_volume(x: u.Q["length"]) -> u.Q["volume"]:
...     return x**3

>>> hessian_cubbe_volume = u.experimental.hessian(cubbe_volume, units=("m",))
>>> hessian_cubbe_volume(u.Q(2.0, "m"))
Quantity(Array(12., dtype=float32...), unit='m')

argnums selects the argument to differentiate with respect to. For f(x, y) = x y**2, the second derivative w.r.t. y is 2 x:

>>> def f(x, y):
...     return x * y**2

>>> hess_y = u.experimental.hessian(f, argnums=1, units=("m", "s"))
>>> hess_y(u.Q(3.0, "m"), u.Q(4.0, "s"))
Quantity(Array(6., dtype=float32...), unit='m')
Parameters:
Return type:

Callable[[Unpack[TypeVarTuple]], TypeVar(R, bound= AbstractQuantity)]

unxt._src.experimental.jacfwd(fun: Callable[[Unpack[Args]], R], argnums: int = 0, *, units: tuple[UnitBase | FunctionUnitBase | str | None, ...])#

Jacobian of fun evaluated column-by-column using forward-mode AD.

In general, if you can use quax.quaxify(jax.jacfwd(func)) (or the syntactic sugar quax.jacfwd(func)), that’s the better option! The difference from those functions is how units are handled. quaxify will directly propagate the units through the automatic differentiation functions. But sometimes that doesn’t work and we need to strip the units and re-apply them. This function does that, using the units kwarg.

See also

jax.jacfwd

The original JAX jacfwd function.

Examples

>>> import jax.numpy as jnp
>>> import unxt as u

>>> def cubbe_volume(x: u.Q["length"]) -> u.Q["volume"]:
...     return x**3

>>> jacfwd_cubbe_volume = u.experimental.jacfwd(cubbe_volume, units=("m",))
>>> jacfwd_cubbe_volume(u.Q(2.0, "m"))
Quantity(Array(12., dtype=float32...), unit='m2')
Parameters:
Return type:

Callable[[Unpack[TypeVarTuple]], TypeVar(R, bound= AbstractQuantity)]

unxt._src.experimental.where(condition: Array | ndarray | bool | number | bool | int | float | complex, x: R, y: AbstractQuantity, /)#

Unit-checked where: both branches must be quantities (experimental).

A strict alternative to jax.numpy.where(). Both x and y must be quantities (a dimensionless Quantity is fine). y is converted to x’s unit – raising if they are not convertible – and the result is returned in x’s unit and concrete type (x typed as R so a checker sees, e.g., Angle in -> Angle out).

Unlike jnp.where, this will not silently reinterpret a raw array as being in the quantity’s unit (see the “jnp.where adopts the quantity’s unit for a raw-array branch” sharp bit): a raw-array branch is rejected, so wrap it as unxt.Quantity(arr, unit) first. jnp.where cannot itself be made strict, because JAX lowers masking ops (triu/tril/trace, where(mask, q, 0)) to the same primitive and relies on the raw zero-fill adopting the unit.

Examples

>>> import jax.numpy as jnp
>>> import unxt as u
>>> from unxt import experimental

>>> cond = jnp.asarray([True, False])

y is converted to x’s unit; the result is in x’s unit:

>>> experimental.where(cond, u.Q([1.0, 2.0], "m"), u.Q([0.003, 0.004], "km"))
Quantity(Array([1., 4.], dtype=float32), unit='m')

Incompatible units raise:

>>> try:
...     experimental.where(cond, u.Q([1.0, 2.0], "m"), u.Q([1.0, 2.0], "s"))
... except Exception as e:
...     print(type(e).__name__)
UnitConversionError

A raw-array branch is rejected rather than silently adopting the unit:

>>> try:
...     experimental.where(cond, u.Q([1.0, 2.0], "m"), jnp.asarray([3.0, 4.0]))
... except TypeError as e:
...     print(e)
unxt.experimental.where requires both branches to be Quantities; ...
Parameters:
Return type:

TypeVar(R, bound= AbstractQuantity)