Penzai - boîte à outils JAX pour construire, modifier et visualiser des réseaux de neurones

• Bibliothèque JAX annoncée par Google DeepMind, conçue pour écrire des modèles sous forme de structures de données Pytree fonctionnelles et faciles à lire
• Met l’accent sur la possibilité de manipuler facilement les modèles même après leur entraînement
  • Adaptée à la recherche liée à la rétro-ingénierie ou à la suppression de composants du modèle, ainsi qu’à l’inspection et à l’analyse des activations internes, à la chirurgie de modèles, au débogage d’architecture, etc.
  • Fournit des outils de visualisation, de modification et d’analyse des modèles, et peut donc aussi être utilisée simplement pour créer et entraîner des modèles
• Se compose d’un ensemble d’outils modulaires conçus ensemble, mais pouvant chacun être utilisés indépendamment

Outils inclus

penzai.nn (pz.nn)

• Bibliothèque de réseaux de neurones déclarative basée sur des combinateurs, alternative à d’autres bibliothèques comme Flax, Haiku, Keras et Equinox
• Expose la structure complète de la passe forward d’un modèle dans le pytree du modèle, ce qui permet de voir toutes les opérations du modèle via un pretty-printing et d’injecter une nouvelle logique d’exécution avec jax.tree_util
• En interne, un modèle n’est rien d’autre qu’un pytree appelable

penzai.treescope (pz.ts)

• Puissant pretty-printer Python interactif fonctionnant comme remplacement drop-in du moteur de rendu IPython/Colab standard
• Conçu pour aider à comprendre les modèles Penzai et d’autres pytrees JAX profondément imbriqués, avec prise en charge native de la visualisation de NDArray de dimension arbitraire

penzai.core.selectors (pz.select)

• Outil polyvalent pour les pytrees qui généralise la syntaxe .at[...].set(...) de JAX à des parcours de pytree arbitraires basés sur les types
• Permet d’effectuer facilement des réécritures complexes ou des correctifs à la volée sur les modèles Penzai et d’autres structures de données

penzai.core.named_axes (pz.nx)

• Système léger d’axes nommés qui vectorise les fonctions JAX ordinaires sur des axes nommés et permet de passer en douceur entre un style de programmation à axes nommés et un style positionnel, sans devoir apprendre une nouvelle API de tableaux

penzai.data_effects (pz.de)

• Système opt-in pour les arguments auxiliaires, les nombres aléatoires et les variables d’état, construit sur le parcours de pytrees et gardant le contrôle sans gêner l’écriture ou l’utilisation des modèles

[brainer]

Y en aurait-il aussi un similaire pour PyTorch ?