Penzai – JAX-Toolkit zum Erstellen, Bearbeiten und Visualisieren von neuronalen Netzen

• Eine von Google DeepMind veröffentlichte JAX-Bibliothek, um Modelle als gut lesbare, funktionale Pytree-Datenstrukturen zu schreiben
• Der Fokus liegt darauf, dass sich auch nach dem Training eines Modells noch leicht damit arbeiten lässt
  • Geeignet für Forschung rund um das Reverse Engineering oder Entfernen von Modellkomponenten, die Inspektion und Untersuchung interner Aktivierungen, Model Surgery, Architektur-Debugging usw.
  • Bietet Werkzeuge zur Visualisierung, Modifikation und Analyse von Modellen und kann daher auch einfach zum Erstellen und Trainieren von Modellen verwendet werden
• Besteht aus einer Sammlung modularer Werkzeuge, die gemeinsam entworfen wurden, aber jeweils unabhängig voneinander verwendet werden können

Enthaltene Werkzeuge

penzai.nn (pz.nn)

• Eine deklarative, kombinatorbasierte Bibliothek für neuronale Netze als Alternative zu anderen Bibliotheken für neuronale Netze wie Flax, Haiku, Keras und Equinox
• Legt die vollständige Struktur des Forward Pass eines Modells im Modell-Pytree offen, sodass man mit Pretty Printing alle Operationen des Modells sehen und mit jax.tree_util neue Runtime-Logik einfügen kann
• Modelle sind intern nichts weiter als aufrufbare Pytrees

penzai.treescope (pz.ts)

• Ein leistungsstarker interaktiver Python-Pretty-Printer, der als Drop-in-Ersatz für den üblichen IPython-/Colab-Renderer funktioniert
• Entwickelt, um das Verständnis von Penzai-Modellen und anderen tief verschachtelten JAX-Pytrees zu erleichtern, mit nativer Unterstützung für die Visualisierung von NDArrays beliebiger Dimension

penzai.core.selectors (pz.select)

• Ein multifunktionales Werkzeug für Pytrees, das die JAX-Syntax .at[...].set(...) auf beliebige typbasierte Pytree-Traversierungen verallgemeinert
• Ermöglicht komplexes Umschreiben oder On-the-fly-Patching von Penzai-Modellen und anderen Datenstrukturen auf einfache Weise

penzai.core.named_axes (pz.nx)

• Ein leichtgewichtiges System für benannte Achsen, das allgemeine JAX-Funktionen über benannte Achsen vektorisieren kann und einen nahtlosen Wechsel zwischen benanntem und positionsbasiertem Programmierstil ermöglicht, ohne dass eine neue Array-API gelernt werden muss

penzai.data_effects (pz.de)

• Ein Opt-in-System für Zusatzargumente, Zufallszahlen und Zustandsvariablen, das auf Pytree-Traversierung basiert und Kontrolle bietet, ohne das Schreiben oder Verwenden von Modellen zu behindern

[brainer]

Gibt es vielleicht auch etwas Ähnliches für Pytorch?