【Hacker News搬运】奔仔：用于构建、编辑和可视化神经网络的JAX研究工具包
-
Title: Penzai: JAX research toolkit for building, editing, and visualizing neural nets
奔仔：用于构建、编辑和可视化神经网络的JAX研究工具包
Text:
Url: https://github.com/google-deepmind/penzai
Penzai 是一个基于 JAX 的研究工具包，用于构建、编辑和可视化神经网络。它专注于简化对训练后模型的操作，例如反向工程、组件消融、内部激活检查、模型手术、架构调试等。Penzai 还提供了可视化、修改和分析神经网络的工具。 Penzai 结构化为一系列模块化工具，每个工具都可以独立使用。主要工具包括： - `penzai.nn` (`pz.nn`)：一个声明式组合器基础的神经网络库，与其他神经网络库如 Flax、Haiku、Keras 或 Equinox 相比，它暴露了模型前向传播的完整结构，使得可以通过漂亮的打印来查看模型所做的所有操作，并通过 `jax.tree_util` 注入新的运行时逻辑。 - `penzai.treescope` (`pz.ts`)：一个超级强大的交互式 Python 漂亮打印器，作为普通 IPython/Colab 渲染器的替换品。它设计用于帮助理解 Penzai 模型和其他深度嵌套的 JAX pytree，并内置了对任意维度的 NDArray 的支持。 - `penzai.core.selectors` (`pz.select`)：一个 pytree 瑞士军刀，推广了 JAX 的 `.at[...].set(...)` 语法到任意类型驱动的 pytree 遍历，并使复杂的重写或现场修补 Penzai 模型和其他数据结构变得容易。 - `penzai.core.named_axes` (`pz.nx`)：一个轻量级的命名轴系统，它将普通的 JAX 函数提升到命名轴上进行向量化，并允许您无缝地在命名和位置编程风格之间切换，而无需学习新的数组 API。 - `penzai.data_effects` (`pz.de`)：一个基于 pytree 遍历的副作用系统，用于侧边参数、随机数和状态变量，它让您掌控一切，而不会妨碍编写或使用您的模型。 Penzai 的文档可以在 [https://penzai.readthedocs.io](https://penzai.readthedocs.io/) 找到。要开始使用 Penzai，您需要首先安装 JAX，然后可以使用以下命令安装 Penzai： ```shell pip install penzai在 Colab 或 IPython 笔记本中工作时，建议将 Penzai 配置为默认的漂亮打印器，并启用一些用于交互式使用的工具：
pz.ts.register_as_default() pz.ts.register_autovisualize_magic() pz.enable_interactive_context()此外，Penzai 还提供了一个简单的神经网络示例，展示了如何初始化参数、捕获和提取元素非线性激活后的激活值。要了解如何使用 Penzai 构建和操作神经网络，建议从 "如何思考在 Penzai 中" 教程开始，或者查看 Penzai 文档中的其他教程。
请注意，Penzai 不是一个正式支持的 Google 产品。
## Post by: mccoyb ### Comments: **yklcs**: I like JAX, and find most of the core functionality as an "accelerated NumPy" great. Ecosystem fragmentation and difficulties in interop make adopting JAX hard though.<p>There's too much fragmentation within the JAX NN library space, which penzai isn't helping with. I wish everyone using JAX could agree on a single set of libraries for NN, optimization, and data loading.<p>PyTorch code can't be called, meaning a lot of reimplementation in JAX is needed when extending and iterating on prior works, which is the case for most of research. Custom CUDA kernels are a bit fiddly too, I haven't been able to bring Gaussian Splatting to JAX yet. > **yklcs**: 我喜欢JAX，并发现大部分核心功能是“；加速NumPy”；太棒了然而，生态系统的碎片化和interop中的困难使得采用JAX变得困难<p> 有；在JAX NN库空间内的碎片太多，而penzai不是；没有帮助。我希望每个使用JAX的人都能就NN、优化和数据加载的一组库达成一致<p> PyTorch代码可以；这意味着在扩展和迭代先前的工作时，需要在JAX中进行大量的重新实现，这是大多数研究的情况。自定义CUDA内核也有点麻烦，我没有；还没有能够将高斯飞溅引入JAX。 **pizza**: I remember pytorch has some pytree capability, no? So is it safe to say that the any-pytree-compatible modules here are already compatible w/ pytorch? > **pizza**: 我记得pytorch有一些pytree功能，不是吗？因此，可以肯定地说，这里的任何pytree兼容模块都已经兼容w；pytorch？ **catgary**: I’ve only been reading through the docs for a few moments, but I’m pleasantly surprised to find they the authors are using effect handlers to handle effectful computations in ML models. I was in the process of translating a model from torch to Jax using Equinox, this makes me think penzai could be a better choice. > **catgary**: 我只阅读了一些文档，但我惊喜地发现，作者正在使用效果处理程序来处理ML模型中的有效计算。我正在使用Equinox将一个模型从torch转换为Jax，这让我觉得penzai可能是一个更好的选择。 **ubj**: Does anyone know if and how well Penzai can work with Diffrax [1]? I currently use Diffrax + Equinox for scientific machine learning. Penzai looks like an interesting alternative to Equinox.<p>[1]: <a href="https://docs.kidger.site/diffrax/" rel="nofollow">https://docs.kidger.site/diffrax/</a> > **ubj**: 有人知道Penzai是否以及如何与Diffrax[1]合作吗？我目前使用Diffrax+Equinox进行科学的机器学习。奔仔看起来像是春分的一个有趣的替代品<p> [1]：<a href=“https://；/；docs.kidger.site#xx2F；diffrax/”rel=“nofollow”>https:///；docs.kidger.site/；diffrax</一 **ein0p**: Looks great, but outside Google I do not personally know anyone who uses Jax, and I work in this space. > **ein0p**: 看起来很棒，但在谷歌之外，我个人不知道有谁使用Jax，我在这个领域工作。