JAX의 새 타입 정의: hijax로 양자화 배열 만들기
Defining new Jax types with hijax
JAX의 hijax는 pytree와 달리 배열의 묶음을 하나의 타입으로 다룰 수 있게 해줍니다. 이 문서는 int8 양자화 배열을 예로 들어, HiType 서브클래스 정의, 값 클래스 등록, primitives 작성, VJP/vmap/샤딩 지원까지 전체 과정을 설명합니다. hijax 타입은 불변식 보존, 명확한 jaxpr 표현, 커스텀 탄젠트 타입 등에서 pytree보다 강력하지만, 아직 실험적 API입니다.
pytree로는 양자화 배열의 탄젠트 타입이 float0 배열로 제한되어 유용한 미분이 불가능하지만, hijax 타입으로는 양자화 값이 근사하는 연속적인 f32 배열을 탄젠트로 선택할 수 있습니다.