Présentation de hijax

JAX utilise des tableaux comme monnaie d'échange. Les fonctions de transformation prennent des tableaux en entrée et produisent des tableaux en sortie. Lorsque vous souhaitez travailler avec des données agrégées, l'outil habituel est un pytree : vous regroupez des tableaux dans des conteneurs, et JAX aplati transparentement le regroupement en ses feuilles de tableau à chaque limite.

Cependant, parfois, la transparence est exactement ce que vous ne voulez pas. Certaines données sont mieux modélisées comme un nouveau type, avec leur propre identité. Les types hijax (ou « hi types ») fournissent cela.

Définition d'un nouveau type

Pour définir un nouveau type, vous devez sous-classer HiType et implémenter les méthodes lo_ty, lower_val et raise_val pour spécifier comment le type et ses valeurs sont convertis en tableaux ordinaires (lojax). Vous devez également enregistrer une classe Python comme portant des valeurs de ce type en utilisant register_hitype.

from jax.experimental.hijax import HiType, ShapedArray, register_hitype
from jax.sharding import NamedSharding

@dataclass(frozen=True)
class QArrayTy(HiType):
  shape: tuple[int, ...]
  sharding: NamedSharding  
  def lo_ty(self):
    scale_sharding = self.sharding.update(spec=jax.P(*self.sharding.spec[:-1]))
    return [ShapedArray(self.shape, jnp.dtype('int8'), sharding=self.sharding),
            ShapedArray(self.shape[:-1], jnp.dtype('float32'), sharding=scale_sharding)]
  def lower_val(self, q):
    return [q.qvalue, q.scale]
  def raise_val(self, qvalue, scale):
    return QArray(qvalue, scale)
register_hitype(QArray, lambda q: QArrayTy(q.qvalue.shape, jax.typeof(q.qvalue).sharding))

Exemple : tableaux quantifiés

Supposons que nous voulions travailler avec des tableaux quantifiés à int8. Un tableau quantifié est en réalité une paire de tableaux : les valeurs int8 et une échelle de point flottant partagée par chaque ligne.

Pour définir ce type, nous créons une classe QArrayTy qui hérite de HiType et implémente les méthodes nécessaires. Nous enregistrons ensuite cette classe comme portant des valeurs de ce type en utilisant register_hitype.

Implications et limites

La définition de nouveaux types JAX avec hijax offre une grande flexibilité pour modéliser des données complexes. Cependant, cela nécessite une bonne compréhension des mécanismes internes de JAX et de la manière dont les types sont représentés et manipulés.

Il est important de noter que hijax est encore expérimental et que les API peuvent évoluer. Il est donc recommandé de consulter la documentation officielle de JAX pour obtenir les dernières informations sur hijax et ses utilisations.