Defining new JAX types with hijax
Summary
This JAX guide introduces hijax types, showing how to define new types by subclassing HiType, implementing lo_ty, lower_val, and raise_val, and registering with register_hitype. It walks through a quantized array example (QArrayTy), demonstrates per-row tangents via to_tangent_aval, and shows how to define hijax primitives with VJP rules. The article also covers how such hi types integrate with autodiff, vmap, scan, and sharding, including explicit mode and mapping specs, and provides additional examples like Rank-1 arrays and Tuples. It emphasizes that hijax is experimental and points to tests and further docs.