penzai

项目介绍: Penzai

Penzai，意为“盆栽”，是一种古老的中国艺术形式，通过将树木和景观置于微型托盘中种植而成，与日本的盆景艺术相似。Penzai项目是一款基于JAX的库，用于将模型以易读、功能性的pytree数据结构编写。它还提供了一系列工具来帮助用户可视化、修改和分析这些模型。

核心特点

Penzai专注于模型训练后的操作，使其成为涉及逆向工程或剔除模型组件的研究的理想选择。用户可以通过它检查和探查模型的内部激活，执行模型手术，调试架构等。当然，如果用户只是想构建和训练模型，Penzai同样可以胜任。

模块化工具集

Penzai由一组可独立使用的模块化工具组成：

互动式Python漂亮打印机：
Treescope是一个超级功能的互动Python美化打印工具，它原为Penzai的一部分，现在作为独立软件包提供。它旨在帮助理解Penzai模型以及其他深层次嵌套的JAX pytrees，内置支持可视化任意维度的NDArrays。
JAX树和数组的操作工具：
- penzai.core.selectors提供一种广义的JAX .at[...].set(...)语法，适合于复杂重写或即兴修补Penzai模型和数据结构。
- penzai.core.named_axes提供一种轻量级命名轴系统，使用户可以轻松在命名和位置编程风格之间切换。
声明式组合器为基础的神经网络库：
- penzai.nn是一个提供模型声明式组合子的库，让用户可以轻松查看和修改模型。模型通过JAX PyTrees表示，可以通过打印查看其执行过程，并使用jax.tree_util注入新的运行时逻辑。
模块化的通用Transformer架构实现：
- penzai.models.transformer是一个参考Transformer实现，可加载Gemma、Llama、Mistral及GPT-NeoX / Pythia架构的预训练权重。通过模块化组件和命名轴，简化了复杂的模型操作流程。

入门指南

在使用Penzai前，用户需要首先安装JAX，然后通过以下命令安装Penzai：

pip install penzai

并导入使用：

import penzai
from penzai import pz

当在Colab或IPython笔记本中工作时，建议配置Treescope为默认美化打印工具，并启用一些用于互动使用的实用程序。

一个简单的使用例子

以下是如何初始化和可视化一个简单神经网络的示例：

from penzai.models import simple_mlp
mlp = simple_mlp.MLP.from_config(
    name="mlp",
    init_base_rng=jax.random.key(0),
    feature_sizes=[8, 32, 32, 8]
)
mlp

用户还可以捕获和提取元素非线性激活后的中间值：

@pz.pytree_dataclass
class AppendIntermediate(pz.nn.Layer):
  saved: pz.StateVariable[list[Any]]
  def __call__(self, x: Any, **unused_side_inputs) -> Any:
    self.saved.value = self.saved.value + [x]
    return x

var = pz.StateVariable(value=[], label="my_intermediates")

saving_model = (
    pz.select(mlp)
    .at_instances_of(pz.nn.Elementwise)
    .insert_after(AppendIntermediate(var))
)

output = saving_model(pz.nx.ones({"features": 8}))
intermediates = var.value