[{"data":1,"prerenderedAt":-1},["ShallowReactive",2],{"similar-google--flax":3,"tool-google--flax":61},[4,18,26,36,44,53],{"id":5,"name":6,"github_repo":7,"description_zh":8,"stars":9,"difficulty_score":10,"last_commit_at":11,"category_tags":12,"status":17},4358,"openclaw","openclaw\u002Fopenclaw","OpenClaw 是一款专为个人打造的本地化 AI 助手，旨在让你在自己的设备上拥有完全可控的智能伙伴。它打破了传统 AI 助手局限于特定网页或应用的束缚，能够直接接入你日常使用的各类通讯渠道，包括微信、WhatsApp、Telegram、Discord、iMessage 等数十种平台。无论你在哪个聊天软件中发送消息，OpenClaw 都能即时响应，甚至支持在 macOS、iOS 和 Android 设备上进行语音交互，并提供实时的画布渲染功能供你操控。\n\n这款工具主要解决了用户对数据隐私、响应速度以及“始终在线”体验的需求。通过将 AI 部署在本地，用户无需依赖云端服务即可享受快速、私密的智能辅助，真正实现了“你的数据，你做主”。其独特的技术亮点在于强大的网关架构，将控制平面与核心助手分离，确保跨平台通信的流畅性与扩展性。\n\nOpenClaw 非常适合希望构建个性化工作流的技术爱好者、开发者，以及注重隐私保护且不愿被单一生态绑定的普通用户。只要具备基础的终端操作能力（支持 macOS、Linux 及 Windows WSL2），即可通过简单的命令行引导完成部署。如果你渴望拥有一个懂你",349277,3,"2026-04-06T06:32:30",[13,14,15,16],"Agent","开发框架","图像","数据工具","ready",{"id":19,"name":20,"github_repo":21,"description_zh":22,"stars":23,"difficulty_score":10,"last_commit_at":24,"category_tags":25,"status":17},3808,"stable-diffusion-webui","AUTOMATIC1111\u002Fstable-diffusion-webui","stable-diffusion-webui 是一个基于 Gradio 构建的网页版操作界面，旨在让用户能够轻松地在本地运行和使用强大的 Stable Diffusion 图像生成模型。它解决了原始模型依赖命令行、操作门槛高且功能分散的痛点，将复杂的 AI 绘图流程整合进一个直观易用的图形化平台。\n\n无论是希望快速上手的普通创作者、需要精细控制画面细节的设计师，还是想要深入探索模型潜力的开发者与研究人员，都能从中获益。其核心亮点在于极高的功能丰富度：不仅支持文生图、图生图、局部重绘（Inpainting）和外绘（Outpainting）等基础模式，还独创了注意力机制调整、提示词矩阵、负向提示词以及“高清修复”等高级功能。此外，它内置了 GFPGAN 和 CodeFormer 等人脸修复工具，支持多种神经网络放大算法，并允许用户通过插件系统无限扩展能力。即使是显存有限的设备，stable-diffusion-webui 也提供了相应的优化选项，让高质量的 AI 艺术创作变得触手可及。",162132,"2026-04-05T11:01:52",[14,15,13],{"id":27,"name":28,"github_repo":29,"description_zh":30,"stars":31,"difficulty_score":32,"last_commit_at":33,"category_tags":34,"status":17},1381,"everything-claude-code","affaan-m\u002Feverything-claude-code","everything-claude-code 是一套专为 AI 编程助手（如 Claude Code、Codex、Cursor 等）打造的高性能优化系统。它不仅仅是一组配置文件，而是一个经过长期实战打磨的完整框架，旨在解决 AI 代理在实际开发中面临的效率低下、记忆丢失、安全隐患及缺乏持续学习能力等核心痛点。\n\n通过引入技能模块化、直觉增强、记忆持久化机制以及内置的安全扫描功能，everything-claude-code 能显著提升 AI 在复杂任务中的表现，帮助开发者构建更稳定、更智能的生产级 AI 代理。其独特的“研究优先”开发理念和针对 Token 消耗的优化策略，使得模型响应更快、成本更低，同时有效防御潜在的攻击向量。\n\n这套工具特别适合软件开发者、AI 研究人员以及希望深度定制 AI 工作流的技术团队使用。无论您是在构建大型代码库，还是需要 AI 协助进行安全审计与自动化测试，everything-claude-code 都能提供强大的底层支持。作为一个曾荣获 Anthropic 黑客大奖的开源项目，它融合了多语言支持与丰富的实战钩子（hooks），让 AI 真正成长为懂上",141543,2,"2026-04-06T11:32:54",[14,13,35],"语言模型",{"id":37,"name":38,"github_repo":39,"description_zh":40,"stars":41,"difficulty_score":32,"last_commit_at":42,"category_tags":43,"status":17},2271,"ComfyUI","Comfy-Org\u002FComfyUI","ComfyUI 是一款功能强大且高度模块化的视觉 AI 引擎，专为设计和执行复杂的 Stable Diffusion 图像生成流程而打造。它摒弃了传统的代码编写模式，采用直观的节点式流程图界面，让用户通过连接不同的功能模块即可构建个性化的生成管线。\n\n这一设计巧妙解决了高级 AI 绘图工作流配置复杂、灵活性不足的痛点。用户无需具备编程背景，也能自由组合模型、调整参数并实时预览效果，轻松实现从基础文生图到多步骤高清修复等各类复杂任务。ComfyUI 拥有极佳的兼容性，不仅支持 Windows、macOS 和 Linux 全平台，还广泛适配 NVIDIA、AMD、Intel 及苹果 Silicon 等多种硬件架构，并率先支持 SDXL、Flux、SD3 等前沿模型。\n\n无论是希望深入探索算法潜力的研究人员和开发者，还是追求极致创作自由度的设计师与资深 AI 绘画爱好者，ComfyUI 都能提供强大的支持。其独特的模块化架构允许社区不断扩展新功能，使其成为当前最灵活、生态最丰富的开源扩散模型工具之一，帮助用户将创意高效转化为现实。",107888,"2026-04-06T11:32:50",[14,15,13],{"id":45,"name":46,"github_repo":47,"description_zh":48,"stars":49,"difficulty_score":32,"last_commit_at":50,"category_tags":51,"status":17},4721,"markitdown","microsoft\u002Fmarkitdown","MarkItDown 是一款由微软 AutoGen 团队打造的轻量级 Python 工具，专为将各类文件高效转换为 Markdown 格式而设计。它支持 PDF、Word、Excel、PPT、图片（含 OCR）、音频（含语音转录）、HTML 乃至 YouTube 链接等多种格式的解析，能够精准提取文档中的标题、列表、表格和链接等关键结构信息。\n\n在人工智能应用日益普及的今天，大语言模型（LLM）虽擅长处理文本，却难以直接读取复杂的二进制办公文档。MarkItDown 恰好解决了这一痛点，它将非结构化或半结构化的文件转化为模型“原生理解”且 Token 效率极高的 Markdown 格式，成为连接本地文件与 AI 分析 pipeline 的理想桥梁。此外，它还提供了 MCP（模型上下文协议）服务器，可无缝集成到 Claude Desktop 等 LLM 应用中。\n\n这款工具特别适合开发者、数据科学家及 AI 研究人员使用，尤其是那些需要构建文档检索增强生成（RAG）系统、进行批量文本分析或希望让 AI 助手直接“阅读”本地文件的用户。虽然生成的内容也具备一定可读性，但其核心优势在于为机器",93400,"2026-04-06T19:52:38",[52,14],"插件",{"id":54,"name":55,"github_repo":56,"description_zh":57,"stars":58,"difficulty_score":10,"last_commit_at":59,"category_tags":60,"status":17},4487,"LLMs-from-scratch","rasbt\u002FLLMs-from-scratch","LLMs-from-scratch 是一个基于 PyTorch 的开源教育项目，旨在引导用户从零开始一步步构建一个类似 ChatGPT 的大型语言模型（LLM）。它不仅是同名技术著作的官方代码库，更提供了一套完整的实践方案，涵盖模型开发、预训练及微调的全过程。\n\n该项目主要解决了大模型领域“黑盒化”的学习痛点。许多开发者虽能调用现成模型，却难以深入理解其内部架构与训练机制。通过亲手编写每一行核心代码，用户能够透彻掌握 Transformer 架构、注意力机制等关键原理，从而真正理解大模型是如何“思考”的。此外，项目还包含了加载大型预训练权重进行微调的代码，帮助用户将理论知识延伸至实际应用。\n\nLLMs-from-scratch 特别适合希望深入底层原理的 AI 开发者、研究人员以及计算机专业的学生。对于不满足于仅使用 API，而是渴望探究模型构建细节的技术人员而言，这是极佳的学习资源。其独特的技术亮点在于“循序渐进”的教学设计：将复杂的系统工程拆解为清晰的步骤，配合详细的图表与示例，让构建一个虽小但功能完备的大模型变得触手可及。无论你是想夯实理论基础，还是为未来研发更大规模的模型做准备",90106,"2026-04-06T11:19:32",[35,15,13,14],{"id":62,"github_repo":63,"name":64,"description_en":65,"description_zh":66,"ai_summary_zh":66,"readme_en":67,"readme_zh":68,"quickstart_zh":69,"use_case_zh":70,"hero_image_url":71,"owner_login":72,"owner_name":73,"owner_avatar_url":74,"owner_bio":75,"owner_company":76,"owner_location":76,"owner_email":77,"owner_twitter":78,"owner_website":79,"owner_url":80,"languages":81,"stars":108,"forks":109,"last_commit_at":110,"license":111,"difficulty_score":10,"env_os":112,"env_gpu":113,"env_ram":112,"env_deps":114,"category_tags":119,"github_topics":120,"view_count":32,"oss_zip_url":76,"oss_zip_packed_at":76,"status":17,"created_at":121,"updated_at":122,"faqs":123,"releases":152},4664,"google\u002Fflax","flax","Flax is a neural network library for JAX that is designed for flexibility.","Flax 是一个专为 JAX 生态打造的高性能神经网络库，其核心设计理念是“灵活性”。它旨在帮助研究人员和开发者摆脱传统框架的束缚，无需等待官方更新功能，只需直接修改训练循环或示例代码，即可轻松尝试全新的训练方法。\n\n针对深度学习模型构建中常见的复杂性与调试难题，Flax 提供了简洁直观的解决方案。特别是其最新推出的 Flax NNX API，创新性地引入了原生 Python 引用语义，允许用户使用标准的 Python 对象来定义模型。这一特性不仅支持可变状态和引用共享，还大幅降低了模型的创建、检查与调试难度，让代码逻辑更符合直觉。\n\nFlax 由 Google Brain 工程师与 JAX 团队紧密协作开发，内置了从线性层、卷积、归一化到注意力机制及循环单元等丰富的常用组件。它非常适合希望深入底层原理、追求极致定制能力的 AI 研究人员和高级开发者。如果你需要在 JAX 上进行前沿探索或快速验证新想法，Flax 将是你得力的科研伙伴。","\u003Cdiv align=\"center\">\n\u003Cimg src=\"https:\u002F\u002Foss.gittoolsai.com\u002Fimages\u002Fgoogle_flax_readme_f9be18cc5c9c.png\" alt=\"logo\">\u003C\u002Fimg>\n\u003C\u002Fdiv>\n\n# Flax: A neural network library and ecosystem for JAX designed for flexibility\n\n[![Flax - Test](https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Factions\u002Fworkflows\u002Fflax_test.yml\u002Fbadge.svg)](https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Factions\u002Fworkflows\u002Fflax_test.yml)\n[![PyPI version](https:\u002F\u002Fimg.shields.io\u002Fpypi\u002Fv\u002Fflax)](https:\u002F\u002Fpypi.org\u002Fproject\u002Fflax\u002F)\n\n[**Overview**](#overview)\n| [**Quick install**](#quick-install)\n| [**What does Flax look like?**](#what-does-flax-look-like)\n| [**Documentation**](https:\u002F\u002Fflax.readthedocs.io\u002F)\n\nReleased in 2024, Flax NNX is a new simplified Flax API that is designed to make\nit easier to create, inspect, debug, and analyze neural networks in\n[JAX](https:\u002F\u002Fjax.readthedocs.io\u002F). It achieves this by adding first class support\nfor Python reference semantics. This allows users to express their models using\nregular Python objects, enabling reference sharing and mutability.\n\nFlax NNX evolved from the [Flax Linen API](https:\u002F\u002Fflax-linen.readthedocs.io\u002F), which\nwas released in 2020 by engineers and researchers at Google Brain in close collaboration\nwith the JAX team.\n\nYou can learn more about Flax NNX on the [dedicated Flax documentation site](https:\u002F\u002Fflax.readthedocs.io\u002F). Make sure you check out:\n\n* [Flax NNX basics](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fnnx_basics.html)\n* [MNIST tutorial](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fmnist_tutorial.html)\n* [Why Flax NNX](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fwhy.html)\n* [Evolution from Flax Linen to Flax NNX](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fguides\u002Flinen_to_nnx.html)\n\n**Note:** Flax Linen's [documentation has its own site](https:\u002F\u002Fflax-linen.readthedocs.io\u002F).\n\nThe Flax team's mission is to serve the growing JAX neural network\nresearch ecosystem - both within Alphabet and with the broader community,\nand to explore the use-cases where JAX shines. We use GitHub for almost\nall of our coordination and planning, as well as where we discuss\nupcoming design changes. We welcome feedback on any of our discussion,\nissue and pull request threads.\n\nYou can make feature requests, let us know what you are working on,\nreport issues, ask questions in our [Flax GitHub discussion\nforum](https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fdiscussions).\n\nWe expect to improve Flax, but we don't anticipate significant\nbreaking changes to the core API. We use [Changelog](https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Ftree\u002Fmain\u002FCHANGELOG.md)\nentries and deprecation warnings when possible.\n\nIn case you want to reach us directly, we're at flax-dev@google.com.\n\n## Overview\n\nFlax is a high-performance neural network library and ecosystem for\nJAX that is **designed for flexibility**:\nTry new forms of training by forking an example and by modifying the training\nloop, not adding features to a framework.\n\nFlax is being developed in close collaboration with the JAX team and\ncomes with everything you need to start your research, including:\n\n* **Neural network API** (`flax.nnx`): Including [`Linear`](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fapi_reference\u002Fflax.nnx\u002Fnn\u002Flinear.html#flax.nnx.Linear), [`Conv`](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fapi_reference\u002Fflax.nnx\u002Fnn\u002Flinear.html#flax.nnx.Conv), [`BatchNorm`](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fapi_reference\u002Fflax.nnx\u002Fnn\u002Fnormalization.html#flax.nnx.BatchNorm), [`LayerNorm`](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fapi_reference\u002Fflax.nnx\u002Fnn\u002Fnormalization.html#flax.nnx.LayerNorm), [`GroupNorm`](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fapi_reference\u002Fflax.nnx\u002Fnn\u002Fnormalization.html#flax.nnx.GroupNorm), [Attention](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fapi_reference\u002Fflax.nnx\u002Fnn\u002Fattention.html) ([`MultiHeadAttention`](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fapi_reference\u002Fflax.nnx\u002Fnn\u002Fattention.html#flax.nnx.MultiHeadAttention)), [`LSTMCell`](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fapi_reference\u002Fflax.nnx\u002Fnn\u002Frecurrent.html#flax.nnx.nn.recurrent.LSTMCell), [`GRUCell`](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fapi_reference\u002Fflax.nnx\u002Fnn\u002Frecurrent.html#flax.nnx.nn.recurrent.GRUCell), [`Dropout`](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fapi_reference\u002Fflax.nnx\u002Fnn\u002Fstochastic.html#flax.nnx.Dropout).\n\n* **Utilities and patterns**: replicated training, serialization and checkpointing, metrics, prefetching on device.\n\n* **Educational examples**: [MNIST](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fmnist_tutorial.html), [Inference\u002Fsampling with the Gemma language model (transformer)](https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Ftree\u002Fmain\u002Fexamples\u002Fgemma).\n\n## Quick install\n\nFlax uses JAX, so do check out [JAX installation instructions on CPUs, GPUs and TPUs](https:\u002F\u002Fjax.readthedocs.io\u002Fen\u002Flatest\u002Finstallation.html).\n\nYou will need Python 3.8 or later. Install Flax from PyPi:\n\n```\npip install flax\n```\n\nTo upgrade to the latest version of Flax, you can use:\n\n```\npip install --upgrade git+https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax.git\n```\n\nTo install some additional dependencies (like `matplotlib`) that are required but not included\nby some dependencies, you can use:\n\n```bash\npip install \"flax[all]\"\n```\n\n## What does Flax look like?\n\nWe provide three examples using the Flax API: a simple multi-layer perceptron, a CNN and an auto-encoder.\n\nTo learn more about the `Module` abstraction, check out our [docs](https:\u002F\u002Fflax.readthedocs.io\u002F), our [broad intro to the Module abstraction](https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fblob\u002Fmain\u002Fdocs\u002Flinen_intro.ipynb). For additional concrete demonstrations of best practices, refer to our\n[guides](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fguides\u002Findex.html) and\n[developer notes](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fdeveloper_notes\u002Findex.html).\n\nExample of an MLP:\n\n```py\nclass MLP(nnx.Module):\n  def __init__(self, din: int, dmid: int, dout: int, *, rngs: nnx.Rngs):\n    self.linear1 = nnx.Linear(din, dmid, rngs=rngs)\n    self.dropout = nnx.Dropout(rate=0.1, rngs=rngs)\n    self.bn = nnx.BatchNorm(dmid, rngs=rngs)\n    self.linear2 = nnx.Linear(dmid, dout, rngs=rngs)\n\n  def __call__(self, x: jax.Array):\n    x = nnx.gelu(self.dropout(self.bn(self.linear1(x))))\n    return self.linear2(x)\n```\n\nExample of a CNN:\n\n```py\nclass CNN(nnx.Module):\n  def __init__(self, *, rngs: nnx.Rngs):\n    self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), rngs=rngs)\n    self.conv2 = nnx.Conv(32, 64, kernel_size=(3, 3), rngs=rngs)\n    self.avg_pool = partial(nnx.avg_pool, window_shape=(2, 2), strides=(2, 2))\n    self.linear1 = nnx.Linear(3136, 256, rngs=rngs)\n    self.linear2 = nnx.Linear(256, 10, rngs=rngs)\n\n  def __call__(self, x):\n    x = self.avg_pool(nnx.relu(self.conv1(x)))\n    x = self.avg_pool(nnx.relu(self.conv2(x)))\n    x = x.reshape(x.shape[0], -1)  # flatten\n    x = nnx.relu(self.linear1(x))\n    x = self.linear2(x)\n    return x\n```\n\nExample of an autoencoder:\n\n\n```py\nEncoder = lambda rngs: nnx.Linear(2, 10, rngs=rngs)\nDecoder = lambda rngs: nnx.Linear(10, 2, rngs=rngs)\n\nclass AutoEncoder(nnx.Module):\n  def __init__(self, rngs):\n    self.encoder = Encoder(rngs)\n    self.decoder = Decoder(rngs)\n\n  def __call__(self, x) -> jax.Array:\n    return self.decoder(self.encoder(x))\n\n  def encode(self, x) -> jax.Array:\n    return self.encoder(x)\n```\n\n## Citing Flax\n\nTo cite this repository:\n\n```\n@software{flax2020github,\n  author = {Jonathan Heek and Anselm Levskaya and Avital Oliver and Marvin Ritter and Bertrand Rondepierre and Andreas Steiner and Marc van {Z}ee},\n  title = {{F}lax: A neural network library and ecosystem for {JAX}},\n  url = {http:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax},\n  version = {0.12.6},\n  year = {2024},\n}\n```\n\nIn the above bibtex entry, names are in alphabetical order, the version number\nis intended to be that from [flax\u002Fversion.py](https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fblob\u002Fmain\u002Fflax\u002Fversion.py), and the year corresponds to the project's open-source release.\n\n## Note\n\nFlax is an open source project maintained by a dedicated team at Google DeepMind, but is not an official Google product.\n\n","\u003Cdiv align=\"center\">\n\u003Cimg src=\"https:\u002F\u002Foss.gittoolsai.com\u002Fimages\u002Fgoogle_flax_readme_f9be18cc5c9c.png\" alt=\"logo\">\u003C\u002Fimg>\n\u003C\u002Fdiv>\n\n# Flax：专为灵活性设计的 JAX 神经网络库及生态系统\n\n[![Flax - 测试](https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Factions\u002Fworkflows\u002Fflax_test.yml\u002Fbadge.svg)](https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Factions\u002Fworkflows\u002Fflax_test.yml)\n[![PyPI 版本](https:\u002F\u002Fimg.shields.io\u002Fpypi\u002Fv\u002Fflax)](https:\u002F\u002Fpypi.org\u002Fproject\u002Fflax\u002F)\n\n[**概述**](#overview)\n| [**快速安装**](#quick-install)\n| [**Flax 是什么样子的？**](#what-does-flax-look-like)\n| [**文档**](https:\u002F\u002Fflax.readthedocs.io\u002F)\n\n2024 年发布的 Flax NNX 是一个简化的全新 Flax API，旨在让在\n[JAX](https:\u002F\u002Fjax.readthedocs.io\u002F) 中创建、检查、调试和分析神经网络变得更加容易。它通过为 Python 引用语义提供一等公民级的支持来实现这一点。这使得用户能够使用常规 Python 对象来表达他们的模型，从而支持引用共享和可变性。\n\nFlax NNX 源自 [Flax Linen API](https:\u002F\u002Fflax-linen.readthedocs.io\u002F)。该 API 于 2020 年由 Google Brain 的工程师和研究人员与 JAX 团队紧密合作发布。\n\n您可以在 [专门的 Flax 文档网站](https:\u002F\u002Fflax.readthedocs.io\u002F) 上了解更多关于 Flax NNX 的信息。请务必查看：\n\n* [Flax NNX 基础](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fnnx_basics.html)\n* [MNIST 教程](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fmnist_tutorial.html)\n* [为什么选择 Flax NNX](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fwhy.html)\n* [从 Flax Linen 到 Flax NNX 的演进](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fguides\u002Flinen_to_nnx.html)\n\n**注意：** Flax Linen 的 [文档有其独立的站点](https:\u002F\u002Fflax-linen.readthedocs.io\u002F)。\n\nFlax 团队的使命是服务于不断增长的 JAX 神经网络研究生态系统——无论是在 Alphabet 内部还是更广泛的社区中——并探索 JAX 发挥优势的用例。我们几乎所有的协调与规划工作都在 GitHub 上进行，同时也在那里讨论即将发生的设计变更。我们欢迎对任何讨论、问题或拉取请求线程提出反馈意见。\n\n您可以在我们的 [Flax GitHub 讨论区](https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fdiscussions) 提出功能请求、告知我们您的工作进展、报告问题或提问。\n\n我们预计会持续改进 Flax，但不打算对核心 API 进行重大破坏性更改。我们会尽可能使用 [Changelog](https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Ftree\u002Fmain\u002FCHANGELOG.md) 条目和弃用警告。\n\n如果您想直接联系我们，我们的邮箱是 flax-dev@google.com。\n\n## 概述\n\nFlax 是一个面向 JAX 的高性能神经网络库及生态系统，其设计宗旨就是 **灵活性**：\n尝试新的训练方式，只需分叉示例并修改训练循环，而无需向框架添加新特性。\n\nFlax 在与 JAX 团队的密切合作下开发，自带您开展研究所需的一切，包括：\n\n* **神经网络 API** (`flax.nnx`)：包含 [`Linear`](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fapi_reference\u002Fflax.nnx\u002Fnn\u002Flinear.html#flax.nnx.Linear)、[`Conv`](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fapi_reference\u002Fflax.nnx\u002Fnn\u002Flinear.html#flax.nnx.Conv)、[`BatchNorm`](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fapi_reference\u002Fflax.nnx\u002Fnn\u002Fnormalization.html#flax.nnx.BatchNorm)、[`LayerNorm`](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fapi_reference\u002Fflax.nnx\u002Fnn\u002Fnormalization.html#flax.nnx.LayerNorm)、[`GroupNorm`](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fapi_reference\u002Fflax.nnx\u002Fnn\u002Fnormalization.html#flax.nnx.GroupNorm)、[Attention](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fapi_reference\u002Fflax.nnx\u002Fnn\u002Fattention.html)（如 [`MultiHeadAttention`](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fapi_reference\u002Fflax.nnx\u002Fnn\u002Fattention.html#flax.nnx.MultiHeadAttention)）、[`LSTMCell`](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fapi_reference\u002Fflax.nnx\u002Fnn\u002Frecurrent.html#flax.nnx.nn.recurrent.LSTMCell)、[`GRUCell`](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fapi_reference\u002Fflax.nnx\u002Fnn\u002Frecurrent.html#flax.nnx.nn.recurrent.GRUCell)、[`Dropout`](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fapi_reference\u002Fflax.nnx\u002Fnn\u002Fstochastic.html#flax.nnx.Dropout)。\n\n* **实用工具与模式**：分布式训练、序列化与检查点保存、指标计算、设备端预取。\n\n* **教学示例**：[MNIST](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fmnist_tutorial.html)、[使用 Gemma 语言模型（Transformer）进行推理\u002F采样](https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Ftree\u002Fmain\u002Fexamples\u002Fgemma)。\n\n## 快速安装\n\nFlax 使用 JAX，因此请先查看 [JAX 在 CPU、GPU 和 TPU 上的安装说明](https:\u002F\u002Fjax.readthedocs.io\u002Fen\u002Flatest\u002Finstallation.html)。\n\n您需要 Python 3.8 或更高版本。从 PyPI 安装 Flax：\n\n```\npip install flax\n```\n\n要升级到最新版本的 Flax，可以使用：\n\n```\npip install --upgrade git+https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax.git\n```\n\n如果某些依赖项需要额外的包（例如 `matplotlib`），而这些包并未包含在内，您可以运行以下命令来安装所有依赖项：\n\n```bash\npip install \"flax[all]\"\n```\n\n## Flax 是什么样子的？\n\n我们提供了三个使用 Flax API 的示例：一个简单的多层感知机、一个卷积神经网络和一个自编码器。\n\n要深入了解 `Module` 抽象，可以查看我们的[文档](https:\u002F\u002Fflax.readthedocs.io\u002F)以及我们关于 `Module` 抽象的[详细介绍](https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fblob\u002Fmain\u002Fdocs\u002Flinen_intro.ipynb)。如需更多最佳实践的具体演示，请参阅我们的\n[指南](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fguides\u002Findex.html)和\n[开发者笔记](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fdeveloper_notes\u002Findex.html)。\n\n多层感知机的示例：\n\n```py\nclass MLP(nnx.Module):\n  def __init__(self, din: int, dmid: int, dout: int, *, rngs: nnx.Rngs):\n    self.linear1 = nnx.Linear(din, dmid, rngs=rngs)\n    self.dropout = nnx.Dropout(rate=0.1, rngs=rngs)\n    self.bn = nnx.BatchNorm(dmid, rngs=rngs)\n    self.linear2 = nnx.Linear(dmid, dout, rngs=rngs)\n\n  def __call__(self, x: jax.Array):\n    x = nnx.gelu(self.dropout(self.bn(self.linear1(x))))\n    return self.linear2(x)\n```\n\n卷积神经网络的示例：\n\n```py\nclass CNN(nnx.Module):\n  def __init__(self, *, rngs: nnx.Rngs):\n    self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), rngs=rngs)\n    self.conv2 = nnx.Conv(32, 64, kernel_size=(3, 3), rngs=rngs)\n    self.avg_pool = partial(nnx.avg_pool, window_shape=(2, 2), strides=(2, 2))\n    self.linear1 = nnx.Linear(3136, 256, rngs=rngs)\n    self.linear2 = nnx.Linear(256, 10, rngs=rngs)\n\n  def __call__(self, x):\n    x = self.avg_pool(nnx.relu(self.conv1(x)))\n    x = self.avg_pool(nnx.relu(self.conv2(x)))\n    x = x.reshape(x.shape[0], -1)  # 展平\n    x = nnx.relu(self.linear1(x))\n    x = self.linear2(x)\n    return x\n```\n\n自编码器的示例：\n\n\n```py\nEncoder = lambda rngs: nnx.Linear(2, 10, rngs=rngs)\nDecoder = lambda rngs: nnx.Linear(10, 2, rngs=rngs)\n\nclass AutoEncoder(nnx.Module):\n  def __init__(self, rngs):\n    self.encoder = Encoder(rngs)\n    self.decoder = Decoder(rngs)\n\n  def __call__(self, x) -> jax.Array:\n    return self.decoder(self.encoder(x))\n\n  def encode(self, x) -> jax.Array:\n    return self.encoder(x)\n```\n\n## 引用 Flax\n\n引用本仓库时，请使用以下格式：\n\n```\n@software{flax2020github,\n  author = {Jonathan Heek and Anselm Levskaya and Avital Oliver and Marvin Ritter and Bertrand Rondepierre and Andreas Steiner and Marc van {Z}ee},\n  title = {{F}lax: A neural network library and ecosystem for {JAX}},\n  url = {http:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax},\n  version = {0.12.6},\n  year = {2024},\n}\n```\n\n在上述 BibTeX 条目中，作者姓名按字母顺序排列；版本号应为 [flax\u002Fversion.py](https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fblob\u002Fmain\u002Fflax\u002Fversion.py) 中的版本号；年份对应于该项目的开源发布年份。\n\n## 注意事项\n\nFlax 是由 Google DeepMind 的一支专业团队维护的开源项目，但它并非 Google 的官方产品。","# Flax 快速上手指南\n\nFlax 是一个为 JAX 设计的高性能神经网络库，主打灵活性。2024 年推出的 **Flax NNX** 是新简化版 API，支持原生 Python 引用语义，让模型定义、调试和分析更加直观。\n\n## 环境准备\n\n在开始之前，请确保满足以下要求：\n\n*   **操作系统**：Linux, macOS, 或 Windows (WSL2 推荐)\n*   **Python 版本**：3.8 或更高\n*   **前置依赖**：必须先安装 **JAX**。\n    *   请访问 [JAX 安装指南](https:\u002F\u002Fjax.readthedocs.io\u002Fen\u002Flatest\u002Finstallation.html) 根据您的硬件（CPU\u002FGPU\u002FTPU）安装对应版本的 JAX。\n    *   *国内用户提示*：如果官方源下载缓慢，可尝试使用清华或阿里镜像源安装 JAX 相关包。\n\n## 安装步骤\n\n### 1. 基础安装\n通过 PyPI 安装稳定版：\n\n```bash\npip install flax\n```\n\n### 2. 安装最新版（可选）\n如需体验最新功能，可从 GitHub 源码安装：\n\n```bash\npip install --upgrade git+https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax.git\n```\n\n### 3. 安装完整依赖（推荐）\n若需使用绘图等额外功能（如 `matplotlib`），建议安装完整包：\n\n```bash\npip install \"flax[all]\"\n```\n\n*国内加速提示*：若 pip 下载速度慢，可添加国内镜像源参数，例如：\n```bash\npip install flax -i https:\u002F\u002Fpypi.tuna.tsinghua.edu.cn\u002Fsimple\n```\n\n## 基本使用\n\nFlax NNX 允许你像编写普通 Python 类一样定义神经网络。以下是一个最简单的多层感知机（MLP）示例：\n\n```py\nimport jax\nimport jax.numpy as jnp\nfrom flax import nnx\n\nclass MLP(nnx.Module):\n  def __init__(self, din: int, dmid: int, dout: int, *, rngs: nnx.Rngs):\n    # 定义网络层\n    self.linear1 = nnx.Linear(din, dmid, rngs=rngs)\n    self.dropout = nnx.Dropout(rate=0.1, rngs=rngs)\n    self.bn = nnx.BatchNorm(dmid, rngs=rngs)\n    self.linear2 = nnx.Linear(dmid, dout, rngs=rngs)\n\n  def __call__(self, x: jax.Array):\n    # 定义前向传播逻辑\n    x = nnx.gelu(self.dropout(self.bn(self.linear1(x))))\n    return self.linear2(x)\n\n# 实例化模型\nrngs = nnx.Rngs(0)\nmodel = MLP(din=10, dmid=20, dout=2, rngs=rngs)\n\n# 创建虚拟输入数据\nx = jnp.ones((4, 10))\n\n# 执行前向推理\noutput = model(x)\nprint(output.shape)  # 输出: (4, 2)\n```\n\n**核心要点：**\n*   继承 `nnx.Module` 类来定义模型。\n*   在 `__init__` 中初始化层，需传入 `rngs` 用于随机数生成。\n*   在 `__call__` 中定义数据流动逻辑。\n*   直接调用模型实例即可进行计算。","某 AI 实验室的研究团队正在基于 JAX 开发一种新型动态图神经网络，需要频繁修改模型内部状态传递逻辑以验证创新算法。\n\n### 没有 flax 时\n- 研究人员被迫使用纯函数式写法手动管理所有模型参数和状态，代码中充斥着繁琐的 `params` 字典传递，难以直观理解网络结构。\n- 每次尝试新的训练循环变体（如梯度裁剪或自定义优化步骤）时，都需要重写大量样板代码，无法灵活复用现有组件。\n- 调试过程极其痛苦，因为模型不可就地修改（immutable），必须追踪复杂的函数返回链才能定位状态更新错误。\n- 实现带有共享权重的复杂拓扑结构时，需手动处理引用关系，极易出错且代码可读性差。\n\n### 使用 flax 后\n- 借助 Flax NNX 的原生 Python 引用语义，研究员可以直接用常规类定义模型，像操作普通对象一样自然地访问和修改层属性。\n- 训练逻辑与模型定义彻底解耦，只需简单叉改示例中的训练循环即可尝试新策略，无需触碰核心模型代码。\n- 支持就地调试和断点检查，开发者能直接观察中间变量变化，快速定位并修复状态更新异常，研发效率显著提升。\n- 轻松实现权重共享等高级架构，通过标准的 Python 对象引用即可让不同模块共享同一层实例，代码简洁且符合直觉。\n\nFlax 将 JAX 的高性能与 Python 的灵活性完美结合，让研究者能专注于算法创新而非框架束缚。","https:\u002F\u002Foss.gittoolsai.com\u002Fimages\u002Fgoogle_flax_f9be18cc.png","google","Google","https:\u002F\u002Foss.gittoolsai.com\u002Favatars\u002Fgoogle_c4bedcda.png","Google ❤️ Open Source",null,"opensource@google.com","GoogleOSS","https:\u002F\u002Fopensource.google\u002F","https:\u002F\u002Fgithub.com\u002Fgoogle",[82,86,90,94,98,102,105],{"name":83,"color":84,"percentage":85},"Jupyter Notebook","#DA5B0B",51.3,{"name":87,"color":88,"percentage":89},"Python","#3572A5",48.3,{"name":91,"color":92,"percentage":93},"C++","#f34b7d",0.2,{"name":95,"color":96,"percentage":97},"Shell","#89e051",0.1,{"name":99,"color":100,"percentage":101},"CMake","#DA3434",0,{"name":103,"color":104,"percentage":101},"Makefile","#427819",{"name":106,"color":107,"percentage":101},"CSS","#663399",7144,795,"2026-04-06T07:35:43","Apache-2.0","未说明","非必需，支持 CPU、GPU 和 TPU（具体取决于 JAX 的安装配置）",{"notes":115,"python":116,"dependencies":117},"Flax 是基于 JAX 的神经网络库，本身不强制要求 GPU，但若要使用 GPU 或 TPU 加速，需参考 JAX 的安装指南进行相应配置（如安装 CUDA\u002FcuDNN 或 TPU 驱动）。可选安装额外依赖（如 matplotlib）可通过 'pip install flax[all]' 实现。","3.8+",[118,64],"jax",[14],[118],"2026-03-27T02:49:30.150509","2026-04-07T07:12:06.632270",[124,129,134,139,144,148],{"id":125,"question_zh":126,"answer_zh":127,"source_url":128},21216,"Flax 的旧版 API (flax.nn) 和新版 Linen API 有什么区别？我应该使用哪个？","Flax 已经进行了重大的 API 重构，推出了新的 Linen API。新版 API 旨在赋予用户更多能力，同时保持简单并完全暴露 JAX 的功能。旧的 `flax.nn` API 已被弃用。所有官方示例已迁移至新 API，大型项目也已通过升级指南完成过渡。建议新用户直接使用 Linen API，老用户参考升级指南进行迁移。如有问题可在讨论板提问。","https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fissues\u002F208",{"id":130,"question_zh":131,"answer_zh":132,"source_url":133},21217,"Flax 返回的参数字典应该是 FrozenDict 还是普通字典？为什么？","Flax 倾向于使用 FrozenDict（不可变字典）来防止意外修改。如果使用普通字典，在函数被 JIT 编译与否的情况下，行为可能会不一致（例如直接修改参数字典可能导致难以调试的错误）。虽然社区有支持普通字典的声音，但为了安全性和一致性，Flax 模块内部通常通过 `apply` 方法的 `mutable` 参数来处理可变状态，而默认返回不可变结构以避免副作用。","https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fissues\u002F1223",{"id":135,"question_zh":136,"answer_zh":137,"source_url":138},21218,"如何在 Flax 中使用 FlashAttention？它会自动启用吗？","如果你使用 `flax.nnx.dot_product_attention` 或 `flax.nnx.MultiHeadAttention`，且在 GPU 上运行，同时设置 `dropout` 为无（或不使用 dropout）且 `sow_weights=False`，那么 Flax 默认调用的 `jax.nn.dot_product_attention` 会自动切换到 `cudnn` 实现，从而使用 Flash Attention 内核。如果你使用的是旧版 `flax.linen` 中的 `nn.MultiHeadDotProductAttention`，则可能不会自动启用，可能需要手动配置或使用 NNX 版本。对于 TPU 用户，`jax.nn.dot_product_attention` 仅提供 XLA 版本，若需更高性能可直接使用 Pallas 版本（API 略有不同，如使用 `segment_id` 和 `causal` 代替 `mask`）。","https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fissues\u002F2858",{"id":140,"question_zh":141,"answer_zh":142,"source_url":143},21219,"为什么在 CPU 上使用不同 batch size 调用 model.apply 时，网络行为不一致（即使使用了固定的 BatchNorm 参数）？","该问题与数据类型（dtype）行为有关。早期版本的 JAX\u002FFlax 在不同 batch size 下可能因默认浮点精度处理不一致导致输出差异。此问题已在后续版本中修复，原因是移除了默认的 float32 dtype 强制行为，使得 dtype 行为现在保持一致。确保你使用的是较新版本的 Flax 和 JAX，并且不要依赖全局默认精度设置（JAX\u002FFlax 不建议这样做，因为会导致库代码隐式依赖全局状态）。如果仍遇到问题，请检查是否显式指定了所有层的 dtype。","https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fissues\u002F1755",{"id":145,"question_zh":146,"answer_zh":147,"source_url":128},21220,"Flax 模块实例化时为什么要重写 __new__？这带来了什么问题？","旧版 Flax 模块重写了 `__new__` 方法，导致模块实例在语义上不直观：看似构造对象（如 `nn.Dense(x, features=10)`）实际上执行了两步操作——先创建实例再调用 `apply`。这导致复用模块时必须使用令人困惑的 `module.shared()` 抽象，且无法在其他模块外直接实例化模块。新版 Linen API 已解决此问题，不再重写 `__new__`，使模块行为更符合标准 Python 对象语义，支持更直观的实例化和复用方式。",{"id":149,"question_zh":150,"answer_zh":151,"source_url":143},21221,"JAX 和 Flax 为什么默认不使用双精度（float64），能否全局设置？","JAX 和 Flax 默认非常保守地使用单精度（float32），不自动启用双精度，也不允许通过全局设置更改默认精度。这是因为全局设置会导致库代码或用户代码隐式依赖该全局状态，从而引发难以追踪的错误和不一致行为。最佳实践是在需要双精度的地方显式指定 dtype（例如 `jnp.array(..., dtype=jnp.float64)`），而不是依赖全局配置。这种设计确保了代码的可移植性和可预测性。",[153,158,163,168,173,178,183,188,193,198,203,208,213,218,223,228,233,238,243,248],{"id":154,"version":155,"summary_zh":156,"released_at":157},127232,"v0.12.6","## 变更内容\n* 在 `decorator_lift_transform` 中为 `wrapped_fn` 的 `self` 参数添加注解，以提高可跟踪性，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5298 中完成。\n* 添加 `with_attributes` 功能，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5308 中完成。\n* 移除 optax 的版本锁定，由 @samanklesaria 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5292 中完成。\n* 移除 nnx.Optimizer 相关的冗余代码检查，由 @vfdev-5 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5226 中完成。\n* 修复当 Generic 作为基类中的最后一个时，PyTreeNode + Generic 会丢失 `__parameters__` 属性的问题，由 @mohsinm-dev 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5237 中完成。\n* 添加 `recursive_map` 示例，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5311 中完成。\n* 在 nnx.eval_shape 中添加分片传播支持（克隆自 #5111），由 @samanklesaria 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5247 中完成。\n* 修复 `eval_shape` 的 `_to_variable` 方法，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5316 中完成。\n* 支持在 `set_metadata` 中使用 `out_sharding` 映射，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5312 中完成。\n* 功能更新 (nnx)：为 `MultiHeadAttention` 添加 GQA 支持，由 @ayulockedin 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5259 中完成。\n* 允许对 `DenyList` 进行比较，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5322 中完成。\n* 为 `jit` 添加 `graph_updates` 参数，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5317 中完成。\n* 为 `shard_map` 添加 `graph_updates` 参数，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5319 中完成。\n* 移除 `jax\u002Ftools\u002Fcolab_tpu.py` 文件，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5324 中完成。\n* 为 `vmap` 添加 `graph_updates` 参数，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5320 中完成。\n* 为 `RngStream` 添加 `split` 方法，由 @samanklesaria 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5270 中完成。\n* 为 `scan` 添加 `graph_updates` 参数，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5327 中完成。\n* 为 `while_loop` 添加 `graph_updates` 参数，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5328 中完成。\n* 添加对 `Average`、`Accuracy` 和 `MultiMetric` 中数据掩码的支持，由 @vfdev-5 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5326 中完成。\n* 为 `fori_loop` 添加 `graph_updates` 参数，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5329 中完成。\n* 为 `pmap`、`grad` 和 `value_and_grad` 添加 `graph_updates` 参数，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5330 中完成。\n* 添加树模式下的 nnx FLIP 功能，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5310 中完成。\n* 项目通过 Copybara 导入：由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5331 中完成。\n* 为 `remat` 添加 `graph_updates` 参数，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5336 中完成。\n* 为 `eval_shape` 和 `checkify` 添加 `graph_updates` 参数，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5338 中完成。\n* 添加兼容性模块，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5340 中完成。\n* 增加更多测试，用于检查变换中的一致别名问题，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5341 中完成。","2026-03-20T21:09:44",{"id":159,"version":160,"summary_zh":161,"released_at":162},127233,"v0.12.5","## 变更内容\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5243 中，由 @copybara-service[bot] 将 data\u002Fstatic 定义拆分并维护在 state 中。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5240 中，由 @copybara-service[bot] 为 nnx.grad 和 nnx.value_and_grad 添加了 graph=False 支持。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5242 中，由 @copybara-service[bot] 为 nnx.remat 添加了树模式支持。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5250 中，由 @copybara-service[bot] 为 nnx.vmap 添加了树模式支持。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5251 中，由 @copybara-service[bot] 添加了树模式参数的错误处理。\n* [flax:benchmarks] 修复 Flax 基准测试。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5241 中完成。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5249 中，由 @copybara-service[bot] 改进了未注册数据的检测。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5258 中，由 @copybara-service[bot] 向 split、state、graphdef 和 clone 添加了 graph 参数。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5261 中，由 @copybara-service[bot] 添加了 nnx_graph_mode 配置标志。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5252 中，由 @copybara-service[bot] 为 nnx.{cond,switch,eval_shape,checkify} 添加了树模式支持。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5262 中，由 @copybara-service[bot] 为 nnx.scan 添加了树模式支持。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5264 中，由 @copybara-service[bot] 为 nnx.custom_vjp 添加了树模式支持。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5263 中，由 @copybara-service[bot] 为 nnx.{while,fori}_loop 添加了树模式支持。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5265 中，由 @copybara-service[bot] 为 pmap 添加了树模式支持。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5266 中，由 @copybara-service[bot] 为 iter_graph 添加了树模式支持。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5267 中，由 @copybara-service[bot] 为 nnx.recursive_map、nnx.view 和 nnx.view_info 添加了树模式支持。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5272 中，由 @copybara-service[bot] 在 ToLinen 变量恢复中使用 linen_vars_to_nnx_attrs。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5268 中，由 @copybara-service[bot] 添加了仅限树模式的 nnx.{vjp,jvp,jit_partial}。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5273 中，由 @copybara-service[bot] 在 ToLinen 变量恢复中使用 linen_vars_to_nnx_attrs。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5282 中，由 @copybara-service[bot] 修复了 nn.vmap 的 _broadcast_prefix_tree。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5286 中，由 @copybara-service[bot] 修复了 nn.vmap 的 split_fn。\n* [pmap] 移除 `jax.config.pmap_shmap_merge`。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5289 中完成。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5269 中，由 @copybara-service[bot] 添加了 set_graph_mode。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5296 中，由 @copybara-service[bot] 为 iter_children 添加了错误处理。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5294 中，由 @cgarciae 公开并更新了 view 指南。\n* 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5295 中，由 @samanklesaria 移除了 chex 依赖。\n* 改进了 tr 的错误信息。","2026-03-04T19:18:32",{"id":164,"version":165,"summary_zh":166,"released_at":167},127234,"v0.12.4","此版本修复了 `nnx.List` 和 `nnx.Sequential` 在被视为 PyTree 时，会将其所有元素作为静态参数传递的问题。\n\n## 变更内容\n* 修复 List 并改进 PyTree：@cgarciae，见 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5072\n* 项目 Copybara 导入：@copybara-service[bot]，见 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5211\n* 修复输入被分片时的掩码分片问题：@vfdev-5，见 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5212\n* nnx.view：@chapman20j，见 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5204\n* bonsai 链接：@chapman20j，见 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5220\n* 为 Embed 层调用添加 `out_sharding` 参数：@samanklesaria，见 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5205\n* 改进 PyTree 的展平与反展平：@copybara-service[bot]，见 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5216\n* 锁定 optax 版本以修复 CI 问题：@samanklesaria，见 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5225\n* 将文档依赖项锁定到较低版本：@samanklesaria，见 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5222\n* 修复 Sequential 对静态属性的解包问题：@samanklesaria，见 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5218\n* 功能（nnx）：添加分组查询注意力（GQA）支持：@ayulockedin，见 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5180\n* 杂项：将 gsutil 的使用迁移到 gcloud storage：@gurusai-voleti，见 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5230\n* 使 `linen.Partitioned` 与 `linen.WeightNorm` 兼容：@copybara-service[bot]，见 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5234\n* 在拆分\u002F状态中维护数据\u002F静态定义：@copybara-service[bot]，见 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5227\n* 添加 `linen.WeightNorm` 和 `linen.with_partitioning` 兼容性的测试：@copybara-service[bot]，见 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5236\n* 添加 Tunix 链接：@copybara-service[bot]，见 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5239\n* 为 nnx.shard_map 添加树模式支持：@copybara-service[bot]，见 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5238\n* 将 NNX Variable 元数据中的 `sharding_names` 重命名为 `out_sharding`：@copybara-service[bot]，见 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5215\n\n## 新贡献者\n* @gurusai-voleti 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5230 中做出了首次贡献\n\n**完整变更日志**：https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fcompare\u002Fv0.12.3...v0.12.4","2026-02-12T19:09:36",{"id":169,"version":170,"summary_zh":171,"released_at":172},127235,"v0.12.3","## 变更内容\n* 设置 numpy\u003C2.4，以修复 CI doctest 中的 DeprecationWarning，由 @vfdev-5 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5163 中提出\n* 忽略 optax 的弃用警告，由 @cgarciae 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5165 中提出\n* 修复通用指南首页，由 @ivrolan 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5139 中提出\n* 移除 linen 到 nnx 教程中 nnx.scan 包装 nnx.split_rngs 的调用，由 @samanklesaria 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5160 中提出\n* 在支持的情况下为线性层添加 out_sharding 参数，由 @jackopenn 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5156 中提出\n* 修复 Python 3.14 的 kw_only_dataclasses（第 2 部分），由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5135 中提出\n* 无公开描述，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5164 中提出\n* 让 _graph_flatten 尊重 nnx.data 声明（扩展 #5140），由 @samanklesaria 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5159 中提出\n* [pmap] 避免在新的 `jax.pmap` 下性能下降，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5152 中提出\n* 支持在解析逻辑规则时处理多个 None 和 UNCONSTRAINED，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5129 中提出\n* 改进 hijax 指南，由 @cgarciae 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5115 中提出\n* 文档：修正拼写错误“paramater”→“parameter”，由 @ayulockedin 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5166 中提出\n* 使 `nnx.pop` 能够移除已播种的属性，由 @samanklesaria 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5133 中提出\n* 将 sharding_names 重命名为 sharding_metadata，由 @samanklesaria 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5089 中提出\n* 修复图开销基准测试中的 bug，由 @samanklesaria 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5183 中提出\n* 使用 antigravity 修复 docstrings 中的拼写错误，由 @vfdev-5 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5145 中提出\n* 文档（nnx）：为 Conv 和 LinearGeneral 添加缺失的函数式参数，由 @ayulockedin 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5174 中提出\n* 空变更，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5184 中提出\n* 在 tabulate 期间使用 nnx split（克隆 #5069），由 @samanklesaria 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5186 中提出\n* 处理 nnx.fori_loop 中的纯主体，由 @samanklesaria 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5141 中提出\n* 修复 _cached_partial 方法中的拼写错误，由 @vfdev-5 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5142 中提出\n* 更新 mnist 示例以使用 NNX（克隆 #5064），由 @samanklesaria 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5188 中提出\n* 文档：修复拼写错误并澄清 Functional API 部分的介绍，由 @Moriyuki-S 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5157 中提出\n* 如果指定了 spmd_axis_name 并应用于 vmapped 轴，则在提升后的 vmap 中对 rngs 进行分片。否则，jax.vmap 会抛出错误：“ValueError: 传递给 vmap 的输入的映射维度应以相同的方式进行分片。收到不一致的轴规范：None vs batch”，这是由于 split_rngs 被复制所致，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5189 中提出\n* 移除由自调用嵌套函数引入的引用循环","2026-01-27T21:03:15",{"id":174,"version":175,"summary_zh":176,"released_at":177},127236,"v0.12.2","## 变更内容\n* [flax:examples:wmt] 小型 linter 修复。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5012 中完成\n* [flax:examples:seq2seq] 基于 seq2seq.ipynb 创建主程序和默认配置。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5119 中完成\n* [flax:examples:vae] 小型 linter 修复。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5014 中完成\n* [flax:examples:gemma] 修复 linter 错误。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5013 中完成\n* [flax:examples:sst2] 修复 pytype 错误。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5118 中完成\n* 允许在 `nnx.PathContains` 中使用子字符串匹配，由 @thijs-vanweezel 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5094 中完成\n* [flax:examples:sst2] 修复笔记本错误。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5122 中完成\n* [flax:examples:ppo] 修复一些 linter 和导入问题。#jax-fixit，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5120 中完成\n* 避免向 `jax.remat` 传递 `concrete` 参数，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5121 中完成\n* [flax:examples:lm1b_nnx] 更新示例以在内部正常运行。#jax-fixit。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5125 中完成\n* [flax:examples:nlp_seq] 创建 main.py 文件，以便使用配置文件运行测试，与其他示例保持一致。#jax-fixit，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5126 中完成\n* [jax:benchmarks] 为几个 flax 示例添加跟踪\u002F降级基准测试。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4911 中完成\n* 从 nnx.jit 中移除 abstracted_axes，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5132 中完成\n* 池化操作，由 @jorisSchaller 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5057 中完成\n* 向 flax.nnx.dot_product_attention 添加 is_causal mask 参数，由 @ibbyml 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5093 中完成\n* 为支持它的层的调用方法添加 out_sharding 参数，由 @samanklesaria 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5102 中完成\n* 临时修复 CI 失败问题，由 @vfdev-5 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5144 中完成\n* 新版本 0.12.2，由 @IvyZX 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5149 中发布\n\n## 新贡献者\n* @thijs-vanweezel 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5094 中做出了首次贡献\n* @ibbyml 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5093 中做出了首次贡献\n\n**完整变更日志**：https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fcompare\u002Fv0.12.1...v0.12.2","2025-12-18T22:35:39",{"id":179,"version":180,"summary_zh":181,"released_at":182},127237,"v0.12.1","## 已弃用的功能\n\n### Variable.value\n`Variable.value` 现已弃用。请考虑以下示例：\n\n```python\nimport jax.numpy as jnp\nimport jax\nfrom flax import nnx\n\nmy_param = nnx.Param({'a': 0.0})\n\n@nnx.jit\ndef f(m):\n    m.value['a'] = 1.0\n    return m\n```\n\n运行 `f(my_param)` 的结果是 `Param(value={'a': 0.0})`，而不是以前的 `Param(value={'a': 1.0})`。这是因为获取值参数会返回一个包含 pytree 值（如字典或列表）的副本。相反，应使用 `__setitem__` 方法来更新值：\n\n```python\n@nnx.jit\ndef f(m):\n    m['a'] = 1.0\n    return m\n```\n\n### nnx.Data 和 nnx.Static\n`nnx.Data` 和 `nnx.Static` 注解现已弃用。要创建 `nnx.Pytree` 或 `nnx.Module` 数据类，请使用新的 `nnx.dataclass`，并以 `nnx.data` 和 `nnx.static` 作为字段描述符。\n\n```python\n# 旧方式\n@dataclasses.dataclass\nclass Foo(nnx.Pytree):\n  a: nnx.Data[int]\n  b: nnx.Static[str]\n\n# 新方式\n@nnx.dataclass\nclass Foo(nnx.Pytree):\n  a: int = nnx.data()\n  b: str = nnx.static()\n```\n\n## 拉取请求\n* 清晰化 `*Norm` 层的文档字符串：在 SPMD jit 下，`axis_index_groups` 未被使用。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4940 中提出。\n* 将 `ArrayRef` 的创建移至 `Variable` 创建的末尾，由 @IvyZX 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4980 中提出。\n* 清理与 jax.Ref 相关的名称，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4988 中提出。\n* 向 `nnx.tabulate` 添加 `compute_flops` 和 `compute_vjp_flops` 选项，由 @samanklesaria 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4948 中提出。\n* 修复 `nnx.tabulate` 在处理空字典或 None 值时崩溃的问题（修复 #4889），由 @mohsinm-dev 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4891 中提出。\n* 使 jax.new_ref \u002F jax.Ref 的导入更具未来兼容性，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4986 中提出。\n* 在 `flax.training.common_utils.stack_forest` 中使用 `jnp.stack` 替代 `np.stack`，由 @vfdev-5 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4991 中提出。\n* 修复了损坏的 `nnx.statelib.diff`，由 @vfdev-5 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4992 中提出。\n* 在 NNX 中实现了谱范数，由 @mattbahr 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4623 中提出。\n* 改进 `Variable.{get,set}_metadata`，由 @cgarciae 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4985 中提出。\n* 将 `iter_children` 和 `iter_modules` 移至函数中，由 @samanklesaria 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4961 中提出。\n* 避免在 Python 3.13+ 下安装、导入或测试 tensorflow-text，由 @jburnim 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5001 中提出。\n* 禁止通过 `setattr` 设置元数据，由 @cgarciae 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4993 中提出。\n* 使用 sphinx 6.2+ 用于文档，该版本与 Python 3.13 兼容，由 @jburnim 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5009 中提出。\n* 从热门层中移除了 `kernel_init` 和 `bias_init` 属性，由 @vfdev-5 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4998 中提出。\n* 从 `jax.experimental.enable_x64` 迁移到 `jax.enable_x64`，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F5011 中提出。\n* 添加 Rngs KeylessInitializers","2025-11-19T21:55:54",{"id":184,"version":185,"summary_zh":186,"released_at":187},127238,"v0.12.0","Flax `0.12.0` 包含了许多更新，以及对 NNX API 的一些重要破坏性变更。\n\n## 破坏性变更\n### Pytree 严格属性\n`nnx.Pytree` 及其子类 `nnx.Module` 现在对包含 Array 的属性以及属性状态的更改更加严格。例如，以下代码现在会失败：\n\n```python\nfrom flax import nnx\nimport jax\nimport jax.numpy as jnp\n\nclass Foo(nnx.Module):\n  def __init__(self, use_bias, rngs):\n    self.layers = [  # 错误\n      nnx.Linear(3, 3, rngs=rngs) for _ in range(5)\n    ]\n    self.bias = None # 状态为 static\n    if use_bias:\n      self.bias = nnx.Param(rngs.params.uniform(3,)) # 错误\n```\n\n出现这种情况的原因有两个：\n1. 包含 Array 的 JAX pytree 结构现在必须使用 `nnx.data` 标记。或者，如果容器 pytree 是 `list` 或 `dict`，可以使用 `nnx.List` 或 `nnx.Dict`，它们还允许混合“data”和“static”元素。\n2. 属性将不再自动更改其状态——现在必须显式地使用 `nnx.data` 或 `nnx.static` 来完成。此外，将 Array 或包含 Array 的结构赋值给 static 属性现在也会报错，因为这些属性不会自动变为 data。\n\n要修复上述问题，只需将 `layers` 定义为 `List` Module，它会自动被识别为 data；并在首次赋值时通过 `nnx.data` 显式声明 `bias` 是 data 属性：\n\n```python\nclass Foo(nnx.Module):\n  def __init__(self, use_bias, rngs):\n    self.layers = nnx.List([  # nnx.data 也可以，但推荐使用 List\n      nnx.Linear(3, 3, rngs=rngs) for _ in range(5)\n    ])\n    self.bias = nnx.data(None)\n    if use_bias:\n      self.bias = nnx.Param(rngs.params.uniform(3,))\n```\n\n更多信息请参阅 [Module & Pytree](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fguides\u002Fpytree.html) 指南。\n\n### 迫切分片\n当提供 `sharding_names` 元数据时，变量现在会立即对其值进行分片。这需要一个 mesh，可以通过传递 `mesh` 元数据属性或通过 `jax.set_mesh` 设置全局 mesh 上下文来提供。这使得变量的分片过程简化到构建时：\n\n```python\njax.config.update('jax_num_cpu_devices', 8)\nmesh = jax.make_mesh((2, 4), ('data', 'model'))\n\nwith jax.set_mesh(mesh):\n  variable = nnx.Param(jnp.ones((16, 32)), sharding_names=(None, 'model'))\n  \nprint(variable.value.sharding)\n```\n\n使用 `nnx.with_partitioning` 初始化器装饰器时，也会发生迫切分片，并且会自动扩展到优化器。这意味着模型和优化器都会在构建时就被分片，而无需再使用较为繁琐的 `nnx.get_partition_spec` + `jax.lax.with_sharding_constraint` + `nnx.update` 模式：\n\n```python\nwith jax.set_mesh(mesh):\n  linear = nnx.Linear(\n    in_features=16, out_features=16, use_bias=False,\n    kernel_init=nnx.with_partitioning(\n      nnx.initializers.lecun_normal\n```","2025-09-25T23:58:17",{"id":189,"version":190,"summary_zh":191,"released_at":192},127239,"v0.11.2","## 变更内容\n`nnx.merge` 现在默认不会为传入状态中的 Variables 创建副本，这意味着新的合并结构会持有对传入 Variables 的引用。这使得一些新的模式成为可能，例如现在可以创建具有相同状态但运行时行为不同的模型：\n\n```python\nmodel = SomeModel(...)\n# 创建评估模型\neval_model = nnx.merge(*nnx.split(model))  # 相同的 Variables，不同的结构\neval_model.eval()\n```\n\n`model` 和 `eval_model` 共享相同的 Variables，因此它们会保持同步，但具有不同的运行时行为。这种方式避免了需要在不同运行时模式之间不断切换单个模型的状态，从而减少了出错的可能性或不必要的重新编译。\n\n如果希望保留旧的行为，请使用 `nnx.merge(..., copy=True)`。\n\n### 拉取请求\n* 添加 Rngs 随机辅助函数，由 @cgarciae 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4876 中完成\n* 修复 identity 的重新导出和文档，由 @jlperla 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4850 中完成\n* 修复 ToLinen 文档字符串的返回描述，由 @mohsinm-dev 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4852 中完成\n* 更新文档构建说明并清理未使用的包，由 @IvyZX 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4885 中完成\n* 改进与 dataclasses 相关的文档，由 @IvyZX 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4884 中完成\n* 修复损坏的贡献文档链接，由 @mohsinm-dev 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4855 中完成\n* 内部变更，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4886 中完成\n* 修复 replace_by_pure_dict 中字符串键的保留问题，由 @mohsinm-dev 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4860 中完成\n* 移除 Conv 和 ConvTranspose 对精确批大小的依赖，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4877 中完成\n* 在 flax 的 traceback_util.register_exclusion 中调用 jax 的 source_info_util.register_exclusion，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4887 中完成\n* 修正 nnx.Optimizer 中的拼写错误，由 @codinfox 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4880 中完成\n* 在 docs_nnx 中公开 split_rngs 的文档字符串，由 @vfdev-5 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4846 中完成\n* 将 sentencepiece 版本固定为 0.2.0 以解决问题，由 @IvyZX 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4892 中完成\n* 放宽重复检查，排除 PartitionSpec.UNCONSTRAINED 等非字符串值，因为这些值可以重复出现，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4881 中完成\n* 添加 find_duplicates 函数，由 @cgarciae 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4894 中完成\n* 分片 API 改进（不破坏兼容性），由 @IvyZX 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4893 中完成\n* 记录 jax.random 的简写方法，由 @cgarciae 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4899 中完成\n* Optimiser 已经使用模型实例化 - 05_vae.py，由 @nenuadrian 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4857 中完成\n* 恢复 _check_carry_same_references 中的 is_leaf 逻辑，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4903 中完成\n* 文档修复：移除过时的建议","2025-08-28T17:55:52",{"id":194,"version":195,"summary_zh":196,"released_at":197},127240,"v0.11.1","## 变更内容\n* 由 @SobhanMP 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4796 中将 `Sequential()` 改为恒等操作\n* 由 @IvyZX 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4795 中添加了 JAX\u002FFlax 关键概念文档\n* 由 @cgarciae 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4859 中进行了多项改进\n* 将 `jax.sharding.use_mesh` 替换为 `jax.set_mesh`。`jax.set_mesh` 既可以作为全局设置器使用，也可以作为上下文管理器使用。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4862 中完成\n* 由 @cgarciae 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4863 中对 Pytree 和 ArrayRef 进行了重构\n* 为 object->pytree 的重命名添加了旧属性。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4864 中完成\n* 在 MNIST 教程的 CNN 中添加了 BatchNorm 层，以提高训练稳定性。由 @sanepunk 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4773 中完成\n* 由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4866 中添加了描述\n* 由 @cgarciae 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4869 中实现了字典的 update 和 pop 操作\n* 由 @cgarciae 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4868 中简化了 nnx_basics\n* 由 @cgarciae 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4878 中更新至版本 0.11.1\n\n## 新贡献者\n* @SobhanMP 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4796 中完成了首次贡献\n* @sanepunk 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4773 中完成了首次贡献\n\n**完整变更日志**: https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fcompare\u002Fv0.11.0...v0.11.1","2025-08-08T21:25:29",{"id":199,"version":200,"summary_zh":201,"released_at":202},127241,"v0.11.0","# v0.11.0 - Pytrees、MutableArrays 等！\r\nFlax 的这个版本引入了一些改进，以增强与原生 JAX 的互操作性，并新增对 `jax.experimental.MutableArray` 的支持。更多相关内容即将发布！不过，为了与 JAX 的惯用方式保持一致，也进行了一些破坏性变更。大多数代码应该仍能正常运行，但以下变化与当前行为有所不同：\n\n* 标准层中的 `Rngs`：所有标准层不再持有构造函数中传入的 `rngs` 对象的共享引用，而是保留一个通过 `fork` 方法生成的 `Rngs` 或 `RngStream` 对象副本。这会影响 [在 NNX 转换中使用 Rngs](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fmigrating\u002Fnnx_010_to_nnx_011.html#using-rngs-in-nnx-transforms) 以及 [加载包含 RNG 的检查点](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fmigrating\u002Fnnx_010_to_nnx_011.html#loading-checkpoints-with-rngs)。\n* [优化器更新](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fmigrating\u002Fnnx_010_to_nnx_011.html#optimizer-updates)：为避免引用共享，优化器抽象不再持有对 `model` 的引用，而是要求在调用 `update` 时将 `model` 作为第一个参数传入。\n* 模块作为 Pytree：模块现在已经是 Pytree 了！这可以避免在与原始 JAX 转换进行简单交互时不必要的 `split` 和 `merge` 操作（如果不使用 MutableArrays，状态仍需手动传播，且引用透明性问题依然存在）。这会影响到使用 `jax.tree.*` API 操作 [包含 NNX 对象的 Pytree](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fmigrating\u002Fnnx_010_to_nnx_011.html#pytrees-containing-nnx-objects) 时的行为。\n\n请查看完整的 [NNX 0.10 到 NNX 0.11 迁移指南](https:\u002F\u002Fflax.readthedocs.io\u002Fen\u002Flatest\u002Fmigrating\u002Fnnx_010_to_nnx_011.html#nnx-0-10-to-nnx-0-11)。\n\n在不久的将来，我们将分享更多关于如何借助新的 Pytree 和 MutableArray 支持，直接利用 JAX 转换来使用 NNX 的新方法。敬请期待！\n\n## 变更内容\n* [nnx] 可变数组 p3，由 @cgarciae 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4755 中实现\n* [nnx] 允许在 ToLinen 中调用方法，由 @cgarciae 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4808 中实现\n* 内部变更，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4807 中完成\n* 在 axes_scan 中保留分片信息，由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4806 中实现\n* 去重贡献和理念部分，并将其移至主站点，由 @IvyZX 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4809 中完成\n* 修复 nnx.remat 文档字符串渲染问题，由 @vfdev-5 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4790 中完成\n* 在 Gemma 指南中添加关于模型在 Kaggle 上许可同意的说明，由 @vfdev-5 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4776 中完成\n* [nnx] ToLinen 添加 abstract_init 标志，由 @cgarciae 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4813 中实现\n* 修改 NNX，使其使用 `id(variable)` 而不是将 nnx.Variables 作为字典，由 @divyashreepathihalli 在 https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4814 中完成\n* 允许在 Flax 的初始化和应用中使用 LazyRngs，由 @cop","2025-07-29T21:04:39",{"id":204,"version":205,"summary_zh":206,"released_at":207},127242,"v0.10.7","## What's Changed\r\n* Added identity export from JAX by @jlperla in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4652\r\n* Fixes a bug in type annotations for scope.param (unbox=True should accept callable[..., T | AxisMEtadata[T]] and return T, while unbox=False should always return the same thing as what callable returning. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4727\r\n* fix merge by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4731\r\n* [nnx] make Variable a pytree by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4728\r\n* [nnx] add JitWrapped API by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4699\r\n* Update JAX nightly index usage  by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4733\r\n* [nnx] mutable array p1 by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4715\r\n* add dataclass by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4739\r\n* [flax] unconditionally register nnx.Variable as a pytree by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4748\r\n* Updated version of pre-commit-hooks in .pre-commit-config.yaml by @vfdev-5 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4746\r\n* Fixed docstring visibility for nnx.eval_shape by @vfdev-5 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4747\r\n* Added keep_rngs arg to MHA to optionally store rngs by @vfdev-5 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4749\r\n* MultiHeadAttention only keeps rngs if dropout_rate is positive by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4750\r\n* [nnx] mutable array p2 by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4741\r\n* Add in_kv_features argument to nnx.MultiHeadAttention, addressing https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fissues\u002F4756. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4757\r\n* Fix broken link for Transforms guide by @nireekshak in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4763\r\n* Minor improvements of lm1b_nnx example by @vfdev-5 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4745\r\n* Fix head CI tests by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4764\r\n* Fix typos by @nireekshak in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4725\r\n* Check for leaves of type variablelib.Variable when getting sharding specs. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4769\r\n* Fixes #1925 non-str dict keys not suppoted in module state by @muhrin in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4563\r\n* Modified the Functional API link by @nireekshak in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4767\r\n* Fix hardcoded link to filter guide in docs by @hamogu in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4768\r\n* Fix bad doc links by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4770\r\n* revise axes_scan to flatten argument pytrees only once by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4772\r\n* Simplify ToNNX access of Linen module methods by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4766\r\n* Use `.input_formats` and `.output_formats` in place of `.input_layouts` and `.output_layouts` respectively. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4784\r\n* Exposed OptState in nnx module by @vfdev-5 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4788\r\n* Fixes colab link for nnx docs by @vfdev-5 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4775\r\n* Internal changes by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4786\r\n* Fix typo in Flax `nnx_basics` doc. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4781\r\n* update version to 0.10.7 by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4798\r\n\r\n## New Contributors\r\n* @nireekshak made their first contribution in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4763\r\n* @muhrin made their first contribution in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4563\r\n* @hamogu made their first contribution in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4768\r\n\r\n**Full Changelog**: https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fcompare\u002Fv0.10.6...v0.10.7","2025-07-02T06:09:24",{"id":209,"version":210,"summary_zh":211,"released_at":212},127243,"v0.10.6","## What's Changed\r\n* Sow top activations based on absolute value. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4670\r\n* Add support for layer-specific rope scale factors. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4672\r\n* Automatic model selection for Gemma 3 models. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4671\r\n* Make LoRA's dtype arg useful by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4681\r\n* [NVIDIA] Support FP8 Einsum Op by @kaixih in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4686\r\n* [nnx] remove deprecated APIs by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4627\r\n* Add `attention_bias` parameter to `MultiHeadDotProductAttention`. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4694\r\n* Unit tests for `attention_bias` parameter to `MultiHeadDotProductAttention`. Add parameter to all overloads to make pytype happy. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4702\r\n* Rollback of attention_bias parameter, because the change overrides the attention bias for injected attention functions. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4703\r\n* Add custom einsum op to Einsum() by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4705\r\n* [nnx] refactor GraphDef by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4630\r\n* Make fully replicated array before saving checkpoints for examples that use pmap.  by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4707\r\n* Fix CI by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4716\r\n* remove \"nnx\" collection in ToLinen by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4708\r\n* [nnx] flaxlib types by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4639\r\n* v0.10.6 by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4724\r\n\r\n\r\n**Full Changelog**: https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fcompare\u002Fv0.10.5...v0.10.6","2025-04-23T20:26:32",{"id":214,"version":215,"summary_zh":216,"released_at":217},127244,"v0.10.5","## What's Changed\r\n* [nnx] fix tabulate by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4580\r\n* Refactor bridge.Module tests from `wrappers_test.py` to another file. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4581\r\n* Avoid calls to jnp.shape for non-array inputs. by @jakevdp in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4592\r\n* remove Embed nan casting by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4600\r\n* Add QK Norm. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4594\r\n* Util to let bridge module work with NNX submodules by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4584\r\n* Add configurable Query Pre Attention scalar. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4595\r\n* Make RoPE Base Frequency configurable. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4596\r\n* [nnx] pytrees are graph nodes by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4547\r\n* Add option to load checkpoints with transposed Gating Einsum. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4597\r\n* add top_p sampling in gemma example by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4591\r\n* Fix position and name of Post Attention Norm. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4598\r\n* Add Sow Config to from_params constructor. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4599\r\n* bridge module with linen submodule by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4604\r\n* Dramatically speed up sampling compilation time by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4574\r\n* [nnx] improve grad docs by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4588\r\n* [nnx] add support for standalone Variables by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4606\r\n* add promote_dtype as a config option for multiple layers by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4613\r\n* Copybara import of the project: by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4616\r\n* Fixed typo in `beam_search` loop. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4615\r\n* support swap model params in gemma sampler by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4614\r\n* Allow bridge module to have 'name' field by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4619\r\n* fix performance guide by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4621\r\n* Copybara import of the project: by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4618\r\n* Add REFLECT padding to convolution layer by @sarlinpe in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4553\r\n* fix trace-level detection by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4527\r\n* Add attribute path customization to bridge modules by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4624\r\n* add reprlib max depth flag by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4632\r\n* Allow custom axis metadata annotation during transforms by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4637\r\n* [bridge module] Allow name arg to represent actual submodule path by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4634\r\n* [nnx] improve Variable proxy for binary operations by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4641\r\n* Fix module stack typing annotation. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4633\r\n* Stop passing reduce_axes to jax.grad, jax.vjp, and jax.value_and_grad. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4617\r\n* discord release webhook by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4646\r\n* [nnx] support Array leaves in graph nodes by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4612\r\n* Roll up package jax version and uv.lock by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4648\r\n* Use jax.nn.dot_product_attention when possible by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4649\r\n* Fix flaky vmap test tolerance. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4653\r\n* Test runner ubuntu upgrade 24.04 by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4659\r\n* Fix lazy_init typo by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4657\r\n* deflake a test by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4663\r\n* v0.10.5 by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4656\r\n\r\n## New Contributors\r\n* @sarlinpe made their first contribution in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4553\r\n\r\n**Full Changelog**: https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fcompare\u002Fv0.10.4...v0.10.5","2025-03-31T15:17:48",{"id":219,"version":220,"summary_zh":221,"released_at":222},127245,"v0.10.4","## What's Changed\r\n* update pypi publish by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4538\r\n* [nnx] register_variable_name refactor by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4540\r\n* added support to the accuracy metric for binary classification by @mattbahr in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4536\r\n* [nnx] bridge Module by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4542\r\n* [nnx] copy _var_metadata by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4548\r\n* [bridge] fix unbox logic by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4551\r\n* Add `is_initializing` API by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4550\r\n* [nnx] Add specific model typing for nnx.Optimizer by @marcelroed in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4470\r\n* Add linen metadata conversion to linx by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4552\r\n* [bridge] improve Module context by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4554\r\n* Raise error if user uses 'name' in bridge module setup by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4555\r\n* Add deprecation warning to all `nnx.State` methods by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4561\r\n* [nnx] add shard_map by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4490\r\n* Fix CI breakages from newest jax by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4576\r\n* [bridge] Set _initializing correctly and avoid return RNG states by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4569\r\n* v0.10.4 by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4579\r\n\r\n## New Contributors\r\n* @mattbahr made their first contribution in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4536\r\n* @marcelroed made their first contribution in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4470\r\n\r\n**Full Changelog**: https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fcompare\u002Fv0.10.3...v0.10.4","2025-02-27T00:10:13",{"id":224,"version":225,"summary_zh":226,"released_at":227},127246,"v0.10.3","## What's Changed\r\n* Fix fori_loop and while_loop on multiple modules by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4390\r\n* Upgrade Flax readme to NNX by @8bitmp3 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4386\r\n* [nnx] add performance guide notebook by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4384\r\n* Automated Code Change by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4393\r\n* [nnx] optimize NodeDef.attributes by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4399\r\n* Fixed the broken link in haiku_to_flax.rst file by @tilakrayal in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4402\r\n* [nnx] optimize Variable by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4400\r\n* Update Flax NNX Randomness by @8bitmp3 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4279\r\n* Remove the repeated methods in flax.nnx.Module documentation by @rajasekharporeddy in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4416\r\n* Fixed the broken link in linen_to_nnx.rst by @tilakrayal in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4415\r\n* [nnx] add FlatState by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4410\r\n* Update `async_checkpointer.py` reference by @emmanuel-ferdman in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4385\r\n* Fix multiple links to Orbax documentation by @Matt-Hurd in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4364\r\n* Update links in `Why Flax NNX` documentation by @rajasekharporeddy in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4425\r\n* [nnx] fix transforms guide by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4421\r\n* Add benchmark on state traversal, and a readme by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4428\r\n* Update Flax NNX performance guide by @8bitmp3 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4401\r\n* Create sharding via Partitioned.get_sharding() by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4427\r\n* Update Flax NNX vs JAX Transformations guide by @8bitmp3 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4286\r\n* Upgrade Flax NNX Gemma Sampling Inference doc by @8bitmp3 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4325\r\n* Update NNX `merge` docs in graph.py by @8bitmp3 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4411\r\n* Fix main and add nnx.fori_loop test by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4472\r\n* Upgrade Flax NNX Filters doc by @8bitmp3 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4199\r\n* Makes flax Modules more compatible with IPython auto-reload. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4420\r\n* [nnx] RNN: add broadcast_rngs and state_axes APIs by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4407\r\n* Allow `nnx.bridge.variables.nnx_attrs_to_linen_vars` take `nnx.VariableState` as argument. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4473\r\n* [nnx] add state summaries for print and display by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4438\r\n* Copybara import of the project: by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4475\r\n* [nnx] add state summaries for print and display by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4477\r\n* CI: add scheduled test against nightly JAX releases by @jakevdp in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4478\r\n* CI: pin actions to specific commits by @jakevdp in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4479\r\n* [nnx] fix MultiMetric typing by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4485\r\n* [nnx] fix ToNNX linen_attributes update by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4486\r\n* Remove usages of orbax_utils.save_args_from_target, as this function does nothing (it used to control a checkpointing behavior that has since been optimized away). by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4482\r\n* [nnx] improve Module docs by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4499\r\n* Update einsum layer for Gemma example by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4498\r\n* [nnx] fix fiddle by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4500\r\n* Don't create param in normalization layers instead of create None-value params. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4501\r\n* Rename variable string mapping utils and move them to variableslib by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4503\r\n* fix LoRA initialization error in nnx layer by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4502\r\n* Remove all `Param(None)` lines from NNX by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4504\r\n* make gemma FFW LoRA friendly by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4510\r\n* Add `nnx.Module.perturb` by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4515\r\n* [nnx] add tabulate by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4493\r\n* batch_norm.rst: == should be = by @cool-RR in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4524\r\n* v0.10.3 by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4525\r\n\r\n## New Contributors\r\n* @tilakrayal made their first contribution in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4402\r\n* @emmanuel-ferdman made their first contribution in https:\u002F\u002Fgithub.com\u002Fgoo","2025-02-10T17:33:51",{"id":229,"version":230,"summary_zh":231,"released_at":232},127247,"v0.10.2","## What's Changed\r\n* Add `nnx.fori_loop` by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4353\r\n* Linesearch (and lbfgs) support by @jlperla in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4351\r\n* Upgrade Flax NNX Haiku Linen migration doc by @8bitmp3 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4200\r\n* Fix PRNG handling in `nn.jit` under `nn.scan`. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4359\r\n* support passing arguments directly to the struct.dataclass decorator by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4275\r\n* Avoid assert_array_equal for PRNG keys. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4363\r\n* [nnx] support pure dicts by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4352\r\n* [nnx] add data parallel toy example by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4354\r\n* Add logical axis global context support for NNX by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4350\r\n* [nnx] fix ToLinen kwargs by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4270\r\n* [nnx] use HashableMapping instead of FrozenDict by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4376\r\n* [nnx] fix while_loop\u002Ffori_loop bug when sharing references by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4379\r\n* Add `flax.nnx.eval_shape` docstring by @8bitmp3 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4374\r\n* Setup the flaxlib in C++, using Meson and Nanobind. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4380\r\n* Add `flax.nnx.remat` docstring by @8bitmp3 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4373\r\n* [nnx] add checkify by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4381\r\n* Lint flax.nnx.while_loop docstring by @8bitmp3 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4371\r\n* Lint flax.nnx.fori_loop docstring by @8bitmp3 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4370\r\n* [nnx] add some optimizations to graph.py by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4377\r\n* update version to 0.10.2 by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4387\r\n\r\n## New Contributors\r\n* @jlperla made their first contribution in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4351\r\n\r\n**Full Changelog**: https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fcompare\u002Fv0.10.1...v0.10.2","2024-11-19T19:58:25",{"id":234,"version":235,"summary_zh":236,"released_at":237},127248,"v0.10.1","## What's Changed\r\n* Add Flax NNX GraphDef docstring by @8bitmp3 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4302\r\n* Flesh out the Haiku\u002FFlax guide by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4305\r\n* [nnx] improve mnist tutorial by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4316\r\n* Update Flax Evolution from Linen to NNX guide by @8bitmp3 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4289\r\n* [nnx] try casting integers keys in State.replace_by_pure_dict by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4317\r\n* Fixed nnx examples bad links in the README.md by @vfdev-5 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4282\r\n* Fix philosophy link by @jorisSchaller in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4313\r\n* [nnx] add gemma notebook by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4075\r\n* [nnx] improve init_cache docs by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4291\r\n* remove markdown from section titles by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4322\r\n* Avoid depending on JAX internals, which are about to change. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4326\r\n* Remove outdated compatibility code. by @jakevdp in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4324\r\n* fix ruff complaints by @levskaya in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4331\r\n* Remove GeGLU activation function and golden tests. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4303\r\n* Avoid using float32 in normalization for mean\u002Fvar and scale\u002Fbias parameters when force_float32_reductions=False by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4314\r\n* Avoid assert_array_equal on PRNG keys. by @jakevdp in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4332\r\n* Fix typos in Flax NNX Migrating from Haiku to Flax by @8bitmp3 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4337\r\n* Add API reference for flax.nnx.nn and improve landing page by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4338\r\n* [nnx] improve transforms guide by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4333\r\n* [nnx] cleanup gemma notebook by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4334\r\n* Remove non-lazy RNG compat mode and flag from flax. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4339\r\n* [nnx] fix custom_vjp by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4306\r\n* Define model surgery in docs by @8bitmp3 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4349\r\n* [nnx] update State and variables docstrings by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4346\r\n* Add NNX transforms `nnx.while_loop` and `nnx.switch` by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4343\r\n* update version to v0.10.1 by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4345\r\n\r\n## New Contributors\r\n* @vfdev-5 made their first contribution in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4282\r\n\r\n**Full Changelog**: https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fcompare\u002Fv0.10.0...v0.10.1","2024-10-31T23:09:59",{"id":239,"version":240,"summary_zh":241,"released_at":242},127249,"v0.10.0","## What's Changed\r\n* [nnx] clear nnx basics pip logs by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4149\r\n* Support linen \u003C-> nnx metadata box converging in `nnx.bridge` by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4145\r\n* Add nnx bridge API reference to site by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4158\r\n* [nnx] use jax-style transforms API in nnx_basics by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4155\r\n* [nnx] improve nnx.scan in_axes\u002Fout_axes by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4157\r\n* Support direct quantization for FP8 matmul by @wenscarl in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3922\r\n* Upgrade Flax NNX Model Surgery by @8bitmp3 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4135\r\n* [nnx] add more Variable proxy methods by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4170\r\n* [nnx] disallow Array leaves by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4172\r\n* Internal change by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4176\r\n* [nnx] improve landing page and nnx_basics messaging by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4168\r\n* Fixes a small bug in flax.linen.share_scope, where the scopes of  children of the module being merged that were created before setup(),were not being updated to point to the new scope, and so they would end up staying under the original tree. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4150\r\n* Move all NNX content up a level to be equal with Linen, to make python packaging more consistent. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4177\r\n* Add a guide for `nnx.bridge` by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4171\r\n* [nnx] improve Optimizer metadata propagation by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4180\r\n* [nnx] enable sharding transformation on integer prefixes by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4185\r\n* Support linen.LogicallyPartitioned \u003C-> nnx.Variable by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4161\r\n* Clean up axis hooks in `nnx.Variable` by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4189\r\n* Merge nnx.errors to flax.errors by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4186\r\n* [nnx] optimize jit by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4191\r\n* Split documentation for Linen and NNX by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4192\r\n* Partially revert #4192 which sets back a bunch of previous merged pushes. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4201\r\n* Align bridge variable tree structures by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4194\r\n* [NNX site] Fix landing page and banner phrasing and add examples page by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4202\r\n* shorten banners by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4206\r\n* Add trimmed Linen to NNX guide by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4209\r\n* Minor documentation fixes for AxisMetadata. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4178\r\n* fix tests for numpy 2.0 compatibility by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4215\r\n* Forward all arguments when using nnx.transforms.deprecated.scan as a decorator. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4208\r\n* [nnx] add transforms guide by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4197\r\n* [nnx] fix transforms guide by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4223\r\n* Flax NNX GSPMD guide by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4220\r\n* Update libraries to use JAX's limited (and ill-advised) trace-state-querying APIs rather than depending on JAX's deeper internals, which are about to change. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4225\r\n* [nnx] add Randomness guide by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4216\r\n* Add pure dict conversion util functions to nnx.State. by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4230\r\n* [nnx] Simplify traversal by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4205\r\n* Fix false positive tracer leaks in flax library. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4232\r\n* [nnx] add flaxlib by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4235\r\n* [nnx] improve docs by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4236\r\n* point nnx banner to flax-linen by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4237\r\n* update banners by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4238\r\n* Fix scale dtype and refactor q_dot_dq by @wenscarl in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4229\r\n* update banners by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4241\r\n* Add redirects for Linen guide links in the NNX site scope. by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4242\r\n* Internal change by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4243\r\n* Copybara import of the project: by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4245\r\n* Update Flax NNX Scale Up SPMD guide by @8bitmp3 in ","2024-10-16T23:25:57",{"id":244,"version":245,"summary_zh":246,"released_at":247},127250,"v0.9.0","## What's Changed\r\n* Add NNX surgery guide by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4005\r\n* Port gemma\u002Ftransformer to NNX by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4019\r\n* upgrade python to 3.10 + use pyupgrade by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4038\r\n* [nnx] add Using Filters guide by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4028\r\n* v0.8.6 by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4040\r\n* allow imagenet training profiling to be disabled in config by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4043\r\n* [nnx] LoRAParam inherits from Param by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3988\r\n* [linen] allows multiple compact methods by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3808\r\n* Added support of NANOO fp8. by @wenchenvincent in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3993\r\n* Add functool.wraps() annotation to flax.nn.jit. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4051\r\n* Fix typo in `nnx_basics` doc by @rajasekharporeddy in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4047\r\n* [nnx] fix Variable overloads and add shape\u002Fdtype properties by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4049\r\n* Stop writing msgpack file for new checkpoints and update empty nodes handling so that it no longer depends on this file. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4039\r\n* [nnx] stabilize unsafe_pytree by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4030\r\n* Stop writing msgpack file for new checkpoints and update empty nodes handling so that it no longer depends on this file. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4055\r\n* [NVIDIA] Rename fp8 custom dtype to `fp32_max_grad` by @kaixih in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3984\r\n* [nnx] fix mnist_tutorial colab link by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4063\r\n* [nnx] fix Accuracy on eager mode by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4065\r\n* Update orbax_upgrade_guide.rst for async checkpointing usage examples by @kaushaladiti-2802 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4036\r\n* Re-enable some tests after Python 3.9 is dropped by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4067\r\n* Rename `nnx.compat` to `nnx.bridge` by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4066\r\n* [nnx] improve mnist tutorial by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4070\r\n* Modify Flax checkpointing in preparation for cl\u002F650338576. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4072\r\n* Remove some outdated backward-compatibility code. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4068\r\n* [NVIDIA] Add a user guide for fp8 by @kaixih in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4076\r\n* [nnx] add extract APIs by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4078\r\n* [example]: remove lm1b useless parallism rules by @knightXun in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4077\r\n* [nnx] improve filters guide by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4059\r\n* [nnx] add call by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4004\r\n* Ignore Orbax warning in deprecated `flax.training.checkpoints.py` to unbreak head doctest by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4092\r\n* fix mypy failures due tu numpy update by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4098\r\n* [linen] generalize transform caching by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4057\r\n* [linen] fold rngs on jit to improve caching by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4064\r\n* Add shape-based lazy init to `LinenToNNX` (prev `LinenWrapper`) by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4081\r\n* [nnx] add reseed by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4099\r\n* [nnx] add split\u002Fmerge_inputs by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4084\r\n* Perform shape checks for self.param AFTER unboxing by @danielwatson6 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4079\r\n* fix restore_checkpoint example in docstring by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4101\r\n* [numpy] Fix users of NumPy APIs that are removed in NumPy 2.0. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4104\r\n* set profile_duration_ms = None as in periodic_actions there's default value for both num_profile_steps and profile_duration_ms, and the profile stopping condition is when both num_profile_steps and profile_duration_ms are satisfied, so setting profile_duration_ms=None so that the passed num_profile_steps value gets used by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4096\r\n* [linen] add share_scope by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4102\r\n* Allow metadata pass-through in flax.struct.field by @cool-RR in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4056\r\n* avoid mixing `einsum_dot_general` and `einsum` argument by specifying them explicitly in the caller. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4115\r\n* Add logging to track deprecated codepaths. by @copy","2024-08-27T17:51:17",{"id":249,"version":250,"summary_zh":251,"released_at":252},127251,"v0.8.5","## What's Changed\r\n* v0.8.5 by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3941\r\n* [nnx] improve vmap axis size detection by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3947\r\n* Add direct penzai.treescope support for NNX objects. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3948\r\n* [nnx] fix nnx_basics dependencies by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3942\r\n* Rename all the NNX tests to internal naming & build conventions. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3952\r\n* updated rng guide by @chiamp in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3912\r\n* upgraded haiku guide to include NNX by @chiamp in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3923\r\n* parameterized NNX transforms tests by @chiamp in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3906\r\n* Simplify extended dtypes rules part 1. Start by removing sharding specific rules from EDtypes. This is because we always want to replicate the trailing dims introduced by Edtypes. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3957\r\n* fix HEAD by @chiamp in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3960\r\n* Minor grammar fixes to NNX documentation. by @mcsmart76 in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3953\r\n* Make FlatState a Mapping instead of a dict by @NeilGirdhar in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3928\r\n* Adding Welford metric. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3959\r\n* Modify Welford metric to return mean value. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3970\r\n* [nnx] make State generic by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3964\r\n* updated NNX nn docstrings by @chiamp in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3972\r\n* make flax work with upcoming JAX change to tree_map (being more careful about by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3976\r\n* updated `nnx.module` docstrings by @chiamp in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3966\r\n* updated `nnx.Conv` and `nnx.ConvTranspose` by @chiamp in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3974\r\n* updated `nnx.graph` docstrings by @chiamp in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3958\r\n* * Adds `pmap` and `Pmap`. `static_broadcasted_argnums`, `donate_argnums`, and `global_arg_shapes` are not yet supported. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3978\r\n* Fixes for batch norm docs by @jkarwowski in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3982\r\n* fix deprecation warning by @chiamp in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3981\r\n* updated NNX `rnglib` docstring by @chiamp in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3980\r\n* updated `nnx.training` by @chiamp in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3975\r\n* updated `nnx.variables` docstrings by @chiamp in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3986\r\n* [nnx] vectorize vmap split counts by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3989\r\n* added `wrt` option to `nnx.Optimizer` by @chiamp in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3983\r\n* Added `nnx.graph.iter_children` by @chiamp in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3991\r\n* [nnx] fix vmap by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3995\r\n* Fix head pytest breakage by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4006\r\n* Helper function for loading params from a linen module by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4012\r\n* Port gemma\u002Flayers to NNX by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4013\r\n* [nnx] fix grad by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4007\r\n* [nnx] add PathContains Filter by @cgarciae in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4011\r\n* Support Python 3.9 by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4018\r\n* Port gemma\u002Fmodules to NNX by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4014\r\n* Internal change to fix current head CI by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4017\r\n* Unpin the Orbax pip version. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4024\r\n* Fix Gemma test to unbreak head by @IvyZX in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4025\r\n* Fix pickling of exceptions by @sanderland in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4002\r\n* Call user-defined variable transforms before determining axis size in nn.vmap. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4026\r\n* CI: add test run against oldest supported jax version by @jakevdp in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3996\r\n* Make `force_fp32_for_softmax` arg in `MultiHeadDotProductAttention` useful.  by @copybara-service in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4029\r\n\r\n## New Contributors\r\n* @mcsmart76 made their first contribution in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3953\r\n* @jkarwowski made their first contribution in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F3982\r\n* @sanderland made their first contribution in https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fpull\u002F4002\r\n\r\n**Full Changelog**: https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fflax\u002Fcompare\u002Fv0.8.4...v0.8.5","2024-06-26T09:27:39"]