8 days ago
- Hijax types (hi types) in JAX allow defining custom types beyond arrays for aggregate data with specific invariants and behaviors.
- To define a new hi type, subclass HiType, implement lo_ty, lower_val, and raise_val for lowering to arrays, and associate it with a value class via register_hitype.
- Primitives using hi types are defined as VJPHiPrimitive subclasses with type signatures mentioning the hi type, enforcing invariants through operations.
- Autodiff support requires implementing to_tangent_aval on the type and VJP/JVP rules on primitives, allowing custom tangent types (e.g., float arrays for quantized int8).
- vmap support involves implementing dec_rank and inc_rank on the type, defining a MappingSpec subclass, and batch rules on primitives, with explicit axis_size for hi type arguments.
- Sharding in explicit mode can be integrated by recording sharding data on the type (e.g., NamedSharding), consuming it in lo_ty, and propagating it in primitive typing rules.
- Example: quantized arrays (QArray) represent int8 values with per-row scales, defined as a hi type with primitives for quantization, dequantization, and operations like matmul.
- Hi types appear as single values in jaxprs (e.g., q8[...]), improving readability and ensuring invariants, unlike pytrees which expose individual leaves.
- Additional examples include rank-1 arrays (outer product representation) and generic tuples (TupTy), demonstrating flexibility in component structures and sharding delegation.
- Hi types support transformations like jit, grad, vmap, and scan, with custom handling for leading axes in scan via leading_axis_spec.