Penzai - kit de herramientas de JAX para construir, editar y visualizar redes neuronales

• Biblioteca de JAX presentada por Google DeepMind para escribir modelos como estructuras de datos Pytree funcionales y fáciles de leer
• Se enfoca en facilitar el trabajo con los modelos incluso después de haber sido entrenados
  • Adecuada para investigación relacionada con ingeniería inversa o eliminación de componentes del modelo, además de inspección e investigación de activaciones internas, cirugía de modelos, depuración de arquitectura y más
  • Proporciona herramientas para visualización, modificación y análisis de modelos, por lo que también puede usarse simplemente para crear y entrenar modelos
• Está compuesta por un conjunto de herramientas modulares diseñadas en conjunto, pero que también pueden usarse de forma independiente

Herramientas incluidas

penzai.nn (pz.nn)

• Biblioteca de redes neuronales declarativa basada en combinadores, como alternativa a otras bibliotecas de redes neuronales como Flax, Haiku, Keras y Equinox
• Expone toda la estructura del forward pass del modelo en el pytree del modelo, por lo que se pueden ver todas las operaciones del modelo con pretty printing e inyectar nueva lógica de ejecución con jax.tree_util
• Internamente, el modelo no es más que un pytree invocable

penzai.treescope (pz.ts)

• Potente pretty-printer interactivo de Python que funciona como reemplazo directo del renderizador estándar de IPython/Colab
• Diseñado para ayudar a entender modelos de Penzai y otros pytree de JAX profundamente anidados, con soporte integrado para visualización de NDArray de dimensión arbitraria

penzai.core.selectors (pz.select)

• Herramienta multifunción para pytree que generaliza la sintaxis .at[...].set(...) de JAX a recorridos de pytree arbitrarios basados en tipos
• Permite realizar fácilmente reescrituras complejas o parches sobre la marcha en modelos de Penzai y otras estructuras de datos

penzai.core.named_axes (pz.nx)

• Sistema ligero de ejes nombrados que vectoriza funciones normales de JAX sobre ejes con nombre y permite cambiar fluidamente entre el estilo de programación con nombres y el estilo posicional sin aprender una nueva API de arreglos

penzai.data_effects (pz.de)

• Sistema opcional para argumentos secundarios, números aleatorios y variables de estado, construido sobre recorridos de pytree y que ofrece control sin interferir con la escritura o el uso del modelo

[brainer]

¿Habrá algo similar para PyTorch?