Skip to content

quax.examples.hijaxยค

An array with physical units whose cotangent carries the inverse units. Built on quax.experimental.hijax; see Quax and hijax for when to reach for it, and quax.examples.unitful for the same physics as a plain Quax type.

Minimum JAX 0.10.2 (quax.experimental.hijax.HIJAX_FLOOR)
Import failure ImportError naming the missing hijax objects
Dimensions shared with quax.examples.unitful: import meters, kilograms, seconds, Dimension from there

Supported operationsยค

A rule is registered for each lax primitive below. Any other primitive falls through to quax.Value.materialise, which raises, so an unsupported operation is an error rather than a silent loss of units.

Primitive Operands Units of the result
mul_p two Unitfuls, or a Unitful and an array product of the operands'
add_p two Unitfuls the operands', which must be equal
integer_pow_p one Unitful the operand's, exponents multiplied
reduce_sum_p one Unitful the operand's
broadcast_in_dim_p one Unitful the operand's

jnp.sum, **, * and + therefore work; -, /, <, jnp.where, jnp.copy and jnp.mean raise.

Supported transformationsยค

Transformation Notes
jax.jit Supported. On JAX 0.10.2 and 0.11.0 only, trips JAX_CHECK_TRACER_LEAKS with a false report; see Limitations
jax.grad, jax.vjp Cotangent units are the inverse of the primal's
jax.jvp Tangent units are the primal's
jax.vmap Requires MAPPED as the in_axes/out_axes entry, and axis_size when no plain-array argument is mapped
jax.lax.scan Works with a Unitful carry; no spec needed
jax.shard_map Not supported: UnitfulArrayTy implements no shard, unshard or nospec

Limitationsยค

  • Broadcasting. mul and add broadcast a scalar operand against an array one. Any other broadcast raises TypeError; use broadcast_in_dim explicitly.
  • Transposing a broadcast. broadcast_in_dim has no transpose rule for a broadcast that stretches a size-1 axis, so differentiating through one raises NotImplementedError.
  • JAX_CHECK_TRACER_LEAKS on JAX 0.10.2 and 0.11.0. On those two releases, consuming a hijax value under jit reports a leaked tracer that is not one. It is a JAX regression rather than anything to fix here, reproducible with one HiPrim and no Quax: 0.10.1 and earlier are unaffected, and 0.11.1 fixed it. Quax's test suite disables the check on the affected versions only, so the report is noticed rather than tolerated if it returns.
  • Dimension identity. As in quax.examples.unitful, two Dimension instances with the same name are different dimensions.

The Quax layerยค

Built on quax.experimental.hijax, which supplies the wrapper base class and generates the dispatch rules.

Unitful is a quax.experimental.hijax.HiValue, so it inherits aval, materialise, and the leaf field holding the UnitfulArray.

quax.examples.hijax.Unitful ยค

An array with physical units, backed by a hijax type.

Same idea as quax.examples.unitful.Unitful, different construction: the units live in the JAX type of the leaf rather than in a static pytree field. Two things follow that the pure-Quax version cannot do. The units appear in jaxprs, so a traced program is checked and readable. And a cotangent carries the inverse units, because the hijax type says so -- differentiating a dimensionless loss with respect to a length gives a per-length gradient, where Quax alone must return the primal's own metadata (see Autodiff).

Refuses to quax.Value.materialise, so a primitive with no registered rule is an error rather than a silent unit loss.

Arguments:

  • array: the array to attach units to. An existing UnitfulArray is taken as-is.
  • units: either a single Dimension, or a dict from Dimension to integer exponent -- e.g. {meters: 1, seconds: -2} for an acceleration.
units property ยค

The units of the wrapped array, as a Units tuple.

array property ยค

The wrapped array, with the units dropped.

Applies the unwrap primitive, so this is usable under a trace as well as eagerly.

__init__(array: typing.Any, units: Dimension | dict[Dimension, int] | Units = ()) -> None ยค

quax.examples.hijax.MAPPED = UnitfulArraySpec() module-attribute ยค

How a UnitfulArray is mapped by vmap.

A batch of unitful arrays is one bigger unitful array with the same units, so the leading axis is the only mapping there is and this spec carries no data. Pass it as a vmap in_axes/out_axes entry, or use MAPPED for the Unitful that holds one.

The hijax layerยค

Values of these types are produced and consumed only by the functions below. See the how-to for the two rules that govern using them in traced code.

quax.examples.hijax.UnitfulArray ยค

An array with units, as hijax sees it.

Deliberately not a pytree: JAX flattens a pytree into its leaves before it consults the hijax type registry, so a value that is a pytree can never carry a hijax type. It follows that this class is opaque to jax.tree utilities, to equinox.filter_*, and to anything else that walks pytrees -- which is exactly why Unitful exists to hold one.

Construct these with wrap rather than by calling this class: under a trace, applying the primitive is what keeps the value in the jaxpr, whereas calling the constructor smuggles a tracer into a container JAX believes is a constant.

__annotations_cache__ = {'array': jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex, 'units': tuple[tuple[quax.examples.unitful._core.Dimension, int], ...]} class-attribute ยค

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_fields__ = {'array': Field(name='array',type=jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex,default=<dataclasses._MISSING_TYPE object at 0x7f39b85c4d70>,default_factory=<dataclasses._MISSING_TYPE object at 0x7f39b85c4d70>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({}),kw_only=False,doc=None,_field_type=_FIELD), 'units': Field(name='units',type=tuple[tuple[quax.examples.unitful._core.Dimension, int], ...],default=<dataclasses._MISSING_TYPE object at 0x7f39b85c4d70>,default_factory=<dataclasses._MISSING_TYPE object at 0x7f39b85c4d70>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({}),kw_only=False,doc=None,_field_type=_FIELD)} class-attribute ยค

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_params__ = _DataclassParams(init=True,repr=True,eq=True,order=False,unsafe_hash=False,frozen=True,match_args=True,kw_only=False,slots=False,weakref_slot=False) class-attribute ยค
__dict__ = mappingproxy({'__module__': 'quax.examples.hijax._unitful_array', '__firstlineno__': 126, '__doc__': 'An array with units, as hijax sees it.\n\nDeliberately *not* a pytree: JAX flattens a pytree into its leaves before it\nconsults the hijax type registry, so a value that is a pytree can never\ncarry a hijax type. It follows that this class is opaque to `jax.tree`\nutilities, to `equinox.filter_*`, and to anything else that walks pytrees --\nwhich is exactly why [`Unitful`][quax.examples.hijax.Unitful] exists to hold\none.\n\nConstruct these with [`wrap`][quax.examples.hijax.wrap] rather than by\ncalling this class: under a trace, applying the primitive is what keeps the\nvalue in the jaxpr, whereas calling the constructor smuggles a tracer into a\ncontainer JAX believes is a constant.\n', '__repr__': <function UnitfulArray.__repr__ at 0x7f39832e33d0>, '__annotate_func__': <function UnitfulArray.__annotate__ at 0x7f39832e3480>, '__static_attributes__': (), '__dict__': <attribute '__dict__' of 'UnitfulArray' objects>, '__weakref__': <attribute '__weakref__' of 'UnitfulArray' objects>, '__dataclass_params__': _DataclassParams(init=True,repr=True,eq=True,order=False,unsafe_hash=False,frozen=True,match_args=True,kw_only=False,slots=False,weakref_slot=False), '__annotations_cache__': {'array': jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex, 'units': tuple[tuple[quax.examples.unitful._core.Dimension, int], ...]}, '__dataclass_fields__': {'array': Field(name='array',type=jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex,default=<dataclasses._MISSING_TYPE object at 0x7f39b85c4d70>,default_factory=<dataclasses._MISSING_TYPE object at 0x7f39b85c4d70>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({}),kw_only=False,doc=None,_field_type=_FIELD), 'units': Field(name='units',type=tuple[tuple[quax.examples.unitful._core.Dimension, int], ...],default=<dataclasses._MISSING_TYPE object at 0x7f39b85c4d70>,default_factory=<dataclasses._MISSING_TYPE object at 0x7f39b85c4d70>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({}),kw_only=False,doc=None,_field_type=_FIELD)}, '__replace__': <function _replace at 0x7f39b85ed7a0>, '__hash__': <function UnitfulArray.__hash__ at 0x7f39832e3a00>, '__init__': <function UnitfulArray.__init__ at 0x7f39832e3690>, '__eq__': <function UnitfulArray.__eq__ at 0x7f39832e3530>, '__setattr__': <function UnitfulArray.__setattr__ at 0x7f39832e38a0>, '__delattr__': <function UnitfulArray.__delattr__ at 0x7f39832e3950>, '__match_args__': ('array', 'units')}) class-attribute ยค

Read-only proxy of a mapping.

__doc__ = 'An array with units, as hijax sees it.\n\nDeliberately *not* a pytree: JAX flattens a pytree into its leaves before it\nconsults the hijax type registry, so a value that is a pytree can never\ncarry a hijax type. It follows that this class is opaque to `jax.tree`\nutilities, to `equinox.filter_*`, and to anything else that walks pytrees --\nwhich is exactly why [`Unitful`][quax.examples.hijax.Unitful] exists to hold\none.\n\nConstruct these with [`wrap`][quax.examples.hijax.wrap] rather than by\ncalling this class: under a trace, applying the primitive is what keeps the\nvalue in the jaxpr, whereas calling the constructor smuggles a tracer into a\ncontainer JAX believes is a constant.\n' class-attribute ยค

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__firstlineno__ = 126 class-attribute ยค

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.

int('0b100', base=0) 4

__match_args__ = ('array', 'units') class-attribute ยค

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__module__ = 'quax.examples.hijax._unitful_array' class-attribute ยค

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__static_attributes__ = () class-attribute ยค

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__weakref__ property ยค

list of weak references to the object

__annotate_func__(format) ยค
__delattr__(name) ยค
__eq__(other) ยค
__hash__() ยค
__init__(array: jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex, units: tuple[tuple[quax.examples.unitful._core.Dimension, int], ...]) -> None ยค
__repr__() -> str ยค
__setattr__(name, value) ยค

quax.examples.hijax.UnitfulArrayTy ยค

The hijax type of a UnitfulArray.

Prints in jaxprs as e.g. u[3]{m s^-1}. Equality includes the units, so JAX's own type checking rejects a rule that returns the wrong units.

__annotations_cache__ = {'shape': tuple[int, ...], 'dtype': typing.Any, 'units': tuple[tuple[quax.examples.unitful._core.Dimension, int], ...]} class-attribute ยค

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_fields__ = {'shape': Field(name='shape',type=tuple[int, ...],default=<dataclasses._MISSING_TYPE object at 0x7f39b85c4d70>,default_factory=<dataclasses._MISSING_TYPE object at 0x7f39b85c4d70>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({}),kw_only=False,doc=None,_field_type=_FIELD), 'dtype': Field(name='dtype',type=typing.Any,default=<dataclasses._MISSING_TYPE object at 0x7f39b85c4d70>,default_factory=<dataclasses._MISSING_TYPE object at 0x7f39b85c4d70>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({}),kw_only=False,doc=None,_field_type=_FIELD), 'units': Field(name='units',type=tuple[tuple[quax.examples.unitful._core.Dimension, int], ...],default=<dataclasses._MISSING_TYPE object at 0x7f39b85c4d70>,default_factory=<dataclasses._MISSING_TYPE object at 0x7f39b85c4d70>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({}),kw_only=False,doc=None,_field_type=_FIELD)} class-attribute ยค

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_params__ = _DataclassParams(init=True,repr=True,eq=True,order=False,unsafe_hash=False,frozen=True,match_args=True,kw_only=False,slots=False,weakref_slot=False) class-attribute ยค
__doc__ = "The hijax type of a [`UnitfulArray`][quax.examples.hijax.UnitfulArray].\n\nPrints in jaxprs as e.g. `u[3]{m s^-1}`. Equality includes the units, so\nJAX's own type checking rejects a rule that returns the wrong units.\n" class-attribute ยค

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__firstlineno__ = 169 class-attribute ยค

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.

int('0b100', base=0) 4

__match_args__ = ('shape', 'dtype', 'units') class-attribute ยค

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__module__ = 'quax.examples.hijax._unitful_array' class-attribute ยค

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__static_attributes__ = () class-attribute ยค

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__annotate_func__(format) ยค
__delattr__(name) ยค
__eq__(other) ยค
__hash__() ยค
__init__(shape: tuple[int, ...], dtype: typing.Any, units: tuple[tuple[quax.examples.unitful._core.Dimension, int], ...]) -> None ยค
__repr__(short_dtypes: bool = False, **kwargs: typing.Any) -> str ยค
__setattr__(name, value) ยค
dec_rank(size: int | None, spec: typing.Any) -> UnitfulArrayTy ยค
inc_rank(size: int | None, spec: typing.Any) -> UnitfulArrayTy ยค
leading_axis_spec() -> UnitfulArraySpec ยค
lo_ty() -> list[typing.Any] ยค
lower_val(value: UnitfulArray) -> list[jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex] ยค
raise_val(array: jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex) -> UnitfulArray ยค
str_short(short_dtypes: bool = False, **kwargs: typing.Any) -> str ยค
to_ct_aval() -> UnitfulArrayTy ยค
to_tangent_aval() -> UnitfulArrayTy ยค
vspace_add(x: UnitfulArray, y: UnitfulArray) -> UnitfulArray ยค
vspace_zero() -> UnitfulArray ยค

quax.examples.hijax.UnitfulArraySpec ยค

How a UnitfulArray is mapped by vmap.

A batch of unitful arrays is one bigger unitful array with the same units, so the leading axis is the only mapping there is and this spec carries no data. Pass it as a vmap in_axes/out_axes entry, or use MAPPED for the Unitful that holds one.

__annotations_cache__ = {} class-attribute ยค

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_fields__ = {} class-attribute ยค

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_params__ = _DataclassParams(init=True,repr=True,eq=True,order=False,unsafe_hash=False,frozen=True,match_args=True,kw_only=False,slots=False,weakref_slot=False) class-attribute ยค
__doc__ = 'How a [`UnitfulArray`][quax.examples.hijax.UnitfulArray] is mapped by `vmap`.\n\nA batch of unitful arrays is one bigger unitful array with the same units,\nso the leading axis is the only mapping there is and this spec carries no\ndata.\nPass it as a `vmap` `in_axes`/`out_axes` entry, or use\n[`MAPPED`][quax.examples.hijax.MAPPED] for the `Unitful` that holds one.\n' class-attribute ยค

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__firstlineno__ = 157 class-attribute ยค

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.

int('0b100', base=0) 4

__match_args__ = () class-attribute ยค

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__module__ = 'quax.examples.hijax._unitful_array' class-attribute ยค

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__static_attributes__ = () class-attribute ยค

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__delattr__(name) ยค
__eq__(other) ยค
__hash__() ยค
__init__() -> None ยค
__repr__() ยค
__setattr__(name, value) ยค

Units is the canonical representation of physical units used throughout this module: a tuple of (Dimension, exponent) pairs, sorted by dimension name, with zero exponents dropped. It is a tuple rather than the dict that quax.examples.unitful uses because a HiType must be hashable. to_units builds one from a Dimension or a dict.

Applying the primitivesยค

quax.examples.hijax.wrap(x: jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex, units: Dimension | dict[Dimension, int] | Units = ()) -> UnitfulArray ยค

Attach units to an array, giving a UnitfulArray.

from quax.examples.hijax import wrap
from quax.examples.unitful import meters
import jax.numpy as jnp

wrap(jnp.asarray(2.0), meters)  # UnitfulArray(2.0, m)

quax.examples.hijax.unwrap(q: UnitfulArray) -> jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex ยค

Drop the units from a UnitfulArray.

quax.examples.hijax.add(x: typing.Any, y: typing.Any) -> UnitfulArray ยค

Add two unitful arrays of the same units and shape.

quax.examples.hijax.mul(x: typing.Any, y: typing.Any) -> UnitfulArray ยค

Multiply two unitful arrays, or one and a plain array.

quax.examples.hijax.int_pow(q: UnitfulArray, y: int) -> UnitfulArray ยค

Raise a unitful array to a static integer power.

quax.examples.hijax.sum(q: UnitfulArray, axes: tuple[int, ...] | None = None) -> UnitfulArray ยค

Sum a unitful array over axes, or over every axis if axes is None.

quax.examples.hijax.broadcast_in_dim(q: UnitfulArray, shape: tuple[int, ...], broadcast_dimensions: tuple[int, ...]) -> UnitfulArray ยค

Broadcast a unitful array, as jax.lax.broadcast_in_dim does for arrays.

quax.examples.hijax.to_units(units: Dimension | dict[Dimension, int] | Units) -> tuple[tuple[quax.examples.unitful._core.Dimension, int], ...] ยค

Canonicalise a units specification into a Units tuple.

Accepts a bare Dimension (exponent 1), a {Dimension: exponent} dict, or an already-canonical tuple. Zero exponents are dropped, so m / m is dimensionless rather than m^0.