GitHub 热门项目: penzai

GitHub 项目：penzai 仓库地址：https://github.com/google-deepmind/penzai 星级：1901 | 作者：google-deepmind 项目描述：用于构建、编辑和可视化神经网络的 JAX 研究工具包。 =================================================== 自述文件内容： #盆栽 > **盆（“笔”，托盘） 栽（“栽”，种植）** - *中国古代的成形艺术 树木和风景的缩影，也称为盆景，是盆景的祖先。 日本盆景艺术* Penzai 是一个 JAX 库，用于将模型编写为清晰的函数式 pytree 数据 结构，以及用于可视化、修改和分析它们的工具。 Penzai 专注于**让模型在模型完成后更容易进行操作 经过培训**，使其成为涉及逆向工程的研究的绝佳选择 或消融模型组件，检查和探测内部激活， 执行模型手术、调试架构等。 （但如果你只是 想要构建和训练模型，您也可以这样做！） 使用 Penzai，您的神经网络可能如下所示：  Penzai 的结构是模块化工具的集合，设计在一起，但 每个都可以独立使用： * 一个超强大的交互式 Python 漂亮打印机： * [Treescope](https://treescope.readthedocs.io/en/stable/) (`pz.ts`): 最初是普通 IPython/Colab 渲染器的直接替代品 是 Penzai 的一部分，但现在作为独立包提供。它的设计目的是 帮助理解 Penzai 模型和其他深度嵌套的 JAX pytree， 内置支持可视化任意维 NDArray。 * 一组 JAX 树和数组操作实用程序： * `penzai.core.selectors` (`pz.select`)：一把 pytree 瑞士军刀， 将 JAX 的 `.at[...].set(...)` 语法推广到任意类型驱动 pytree 遍历，并且可以轻松进行复杂的重写或 Penzai 模型和其他数据结构的即时修补。 * `penzai.core.named_axes` (`pz.nx`)：一个轻量级的命名轴系统， 提升普通 JAX 函数以在命名轴上进行矢量化，并允许您 在命名和位置编程风格之间无缝切换，无需 必须学习新的数组 API。 * 一个基于组合器的声明性神经网络库，其中模型是 表示为易于修改的数据结构： * `penzai.nn` (`pz.nn`)：其他神经网络库的替代品，例如 Flax、Haiku、Keras 或 Equinox，它公开了模型的完整结构 使用声明性组合器进行前向传递。像 Equinox 一样，模型是 表示为 JAX PyTrees，这意味着您可以看到模型的所有内容 通过漂亮地打印它，并使用“jax.tree_util”注入新的运行时逻辑来实现。 然而，“penzai.nn”模型的叶子也可能包含可变变量 树的，允许他们跟踪可变状态和参数 分享。 * 通用 Transformer 架构的模块化实现，以支持 可解释性、模型手术和训练动态的研究： * `penzai.models.transformer`：一个参考 Transformer 实现 可以加载 Gemma、Llama、Mistral 和 GPT-Ne