Hijax: Neue JAX-Typen mit eigener Identität definieren

Defining new Jax types with hijax

JAX-Dokumentation stellt Hijax-Typen ("hi types") vor, eine experimentelle Möglichkeit, eigene Typen zu definieren, die als einzelne Werte in jaxprs erscheinen, Invarianten durch primitive Operationen erzwingen und eigene Tangenten- und Batching-Konzepte haben. Am Beispiel eines quantisierten Array-Typs wird gezeigt, wie man HiType-Unterklassen, Primitive mit VJP-Regeln und Unterstützung für vmap und Sharding implementiert. Der Artikel erklärt, warum Pytrees für solche Daten ungeeignet sind und wie Hijax die Kontrolle über Typidentität und Invarianten zurückgibt.