Penzai - kit de ferramentas JAX para construir, editar e visualizar redes neurais

• Biblioteca JAX anunciada pelo Google DeepMind para escrever modelos em uma estrutura de dados Pytree funcional e fácil de ler
• Foco em facilitar o trabalho com o modelo mesmo depois de ele ter sido treinado
  • Adequada para pesquisas relacionadas a engenharia reversa ou remoção de componentes do modelo, inspeção e investigação de ativações internas, cirurgia de modelos, depuração de arquitetura etc.
  • Fornece ferramentas para visualização, modificação e análise de modelos, então também pode ser usada simplesmente para criar e treinar modelos
• Composta por um conjunto de ferramentas modulares projetadas em conjunto, mas que podem ser usadas de forma independente

Ferramentas incluídas

penzai.nn (pz.nn)

• Biblioteca de redes neurais declarativa baseada em combinadores, como alternativa a outras bibliotecas de redes neurais como Flax, Haiku, Keras e Equinox
• Expõe toda a estrutura do forward pass do modelo na pytree do modelo, permitindo ver todas as operações do modelo com pretty printing e injetar nova lógica de runtime com jax.tree_util
• Internamente, o modelo é apenas uma pytree chamável

penzai.treescope (pz.ts)

• Poderoso pretty-printer interativo para Python que funciona como substituição drop-in do renderizador padrão do IPython/Colab
• Projetado para ajudar a entender modelos Penzai e outras pytrees JAX profundamente aninhadas, com suporte nativo à visualização de NDArray de dimensão arbitrária

penzai.core.selectors (pz.select)

• Ferramenta multifuncional para pytrees que generaliza a sintaxe .at[...].set(...) do JAX para travessia de pytree arbitrária baseada em tipos
• Permite reescritas complexas ou patches on-the-fly com facilidade em modelos Penzai e outras estruturas de dados

penzai.core.named_axes (pz.nx)

• Sistema leve de eixos nomeados que vetoriza funções JAX comuns sobre eixos nomeados e permite alternar sem atrito entre o estilo de programação com nomes e o estilo posicional sem aprender uma nova API de arrays

penzai.data_effects (pz.de)

• Sistema opt-in para argumentos auxiliares, números aleatórios e variáveis de estado, construído sobre travessia de pytree e que mantém o controle sem atrapalhar a escrita ou o uso do modelo

[brainer]

Será que existe algo parecido para PyTorch também?