quax.experimental¤

Experimental APIs, which may change or be removed without notice.

quax.experimental.hijax¤

Plumbing for building a quax.Value around a hijax type. Quax and hijax covers when to reach for it, the how-to covers how, and quax.examples.hijax is a worked type built on this module.

Doubly experimental: jax.experimental.hijax is itself experimental and has renamed things in most recent releases. This module resolves those names by probing, so a rename upstream is absorbed here rather than in your code.


Minimum JAX	0.10.2 (`quax.experimental.hijax.HIJAX_FLOOR`)
Import failure	`ImportError` naming the missing hijax objects
Re-exported hijax names	`HiPrim`, `HiType`, `MappingSpec`, `ShapedArray`, `instantiate_zeros`, `register_hitype`

HiPrim is VJPHiPrimitive up to JAX 0.11.1 and HiPrim after it; importing it from here means never spelling either name.

`quax.experimental.hijax.HiValue` ¤

A quax.ArrayValue whose single pytree leaf is a hijax value.

Subclass this, then wire up primitives with register_rules. The subclass normally adds a friendlier __init__ and some read-only properties, and nothing else:

class Unitful(HiValue):
    def __init__(self, array, units=()):
        self.leaf = array if isinstance(array, UnitfulArray) else wrap(array, units)

    @property
    def units(self):
        return jax.typeof(self.leaf).units

Two constraints this class exists to enforce, both of which are easy to get wrong by hand:

aval() reports a plain ShapedArray built from the leaf's hijax type, never the hijax type itself. jax.Array's isinstance check reads a tracer's aval and every jnp function gates on it, so a HiType aval would make jnp reject your tracers -- defeating the point of using Quax.
The leaf must not be re-wrapped. A hijax tracer forwards .shape and .dtype from its type, so it passes equinox.is_array_like; an __init__ that tests for array-likeness will wrap a traced value twice. Test for your hi value class instead.

`init(leaf: typing.Any) -> None` ¤

Arguments:

leaf: the hijax value to wrap.

`aval() -> ShapedArray` ¤

The shape and dtype of the leaf's hijax type, as a ShapedArray.

`materialise() -> typing.Any` ¤

Refuse: materialising would discard whatever the hi type carries.

A primitive with no registered rule is therefore an error rather than a silent downgrade to a plain array. Override if your type would rather fall back than fail.

`quax.experimental.hijax.register_rules(cls: type[quax.experimental.hijax.HiValue], rules: collections.abc.Mapping[jax._src.core.Primitive, collections.abc.Callable[..., typing.Any]]) -> None` ¤

Wire lax primitives to hijax functions for a HiValue subclass.

Each {primitive: function} entry generates the Quax dispatch rules for that primitive: one per combination of cls and plain-array operands, since a binary primitive can be reached with your type on either side or both. Each generated rule unwraps the leaves, applies function, and wraps a hijax result back up in cls.

The number of operands is read from function's positional parameters, and the primitive's own parameters are forwarded to whichever keyword-only arguments function declares. Where the two do not line up -- because the hijax function takes them positionally, or names them differently -- adapt with a small lambda.

Example

register_rules(
    Unitful,
    {
        lax.mul_p: hi.mul,  # (x, y), no params
        lax.add_p: hi.add,
        lax.integer_pow_p: lambda q, *, y, **kw: hi.int_pow(q, y),
        lax.reduce_sum_p: lambda q, *, axes, **kw: hi.sum(q, axes),
    },
)

Arguments:

cls: the HiValue subclass to register against.
rules: a mapping from jax.extend.core.Primitive to the hijax function implementing it.

Returns:

Nothing; the rules are registered as a side effect, exactly as quax.register does.

Raises:

NotImplementedError for a primitive with multiple results, which this helper cannot wrap. Write the rule by hand with quax.register.

quax.experimental¤