JAX v0.5.0发布：随机数生成器改进与重要变更解析

JAX是由Google开发的一个高性能数值计算库，它结合了NumPy的易用性与自动微分、硬件加速等高级功能，特别适合机器学习研究和科学计算。JAX通过XLA编译器将Python/NumPy代码转换为高效的机器代码，可以在CPU、GPU和TPU上运行。

版本变更概述

JAX v0.5.0是一个重要的版本更新，采用了基于工作量(effort-based)的版本控制方案。本次更新包含了一些破坏性变更，特别是随机数生成器(PRNG)的语义变化，可能需要用户更新代码。

主要变更内容

1. 随机数生成器改进

本次版本默认启用了jax_threefry_partitionable特性，这是对JAX随机数生成系统的重要改进。Threefry是一种基于计数器的随机数生成算法，新版本使其支持更好的分区能力，这对于分布式计算场景特别有价值。

2. 平台支持调整

JAX v0.5.0停止了对Mac x86架构的支持，仅保留对Mac ARM架构的支持。这一决策基于两个主要原因：

• Mac x86构建存在多个测试失败和崩溃问题
• Mac x86硬件已停产，开发者难以获取设备进行问题修复

如果社区愿意帮助维护Mac x86平台，特别是确保JAX测试套件在该平台上完全通过，团队会考虑重新支持该架构。

3. 依赖版本要求提升

• 最低NumPy版本要求提升至1.25，该版本将支持到2025年6月
• 最低SciPy版本要求提升至1.11，同样支持到2025年6月

功能改进

1. FFT功能增强

jax.numpy.fft模块中的多维FFT函数(包括fftn、rfftn、ifftn和irfftn)现在支持超过3维的变换，突破了之前的限制。

2. 张量运算优化

jax.numpy.einsum函数的默认优化参数从'optimal'改为'auto'，这避免了在多参数情况下出现指数级增长的计算时间。

3. 线性代数改进

jax.numpy.linalg.solve不再支持右侧参数为批量1D数组的情况。如需保留原有行为，可以使用solve(a, b[..., None]).squeeze(-1)。

新特性

1. FFI用户自定义状态支持

新增了jax.ffi.register_ffi_type_id函数，允许用户在FFI(外部函数接口)中定义自己的状态类型。

2. AOT调试信息支持

AOT(提前编译)降低级的.as_text()方法现在支持debug_info选项，可以在输出中包含调试信息，如源代码位置等。

废弃和移除的功能

1. 废弃的功能

• jax.interpreters.xla模块中的abstractify和pytype_aval_mappings已被废弃，相应功能已移至jax.core模块
• jax.scipy.special.lpmn和jax.scipy.special.lpmn_values函数已被废弃，这与SciPy v1.15.0的变更保持一致
• jax.extend.ffi子模块已移至jax.ffi，原路径将被废弃

2. 移除的功能

• jax_enable_memories标志已被移除，其功能现在默认启用
• jax.lib.xla_client中的Device和XlaRuntimeError已被移除，应使用jax.Device和jax.errors.JaxRuntimeError替代
• jax.experimental.array_api模块已被移除，其功能已直接集成到jax.numpy中

升级建议

对于使用JAX进行科学计算或机器学习开发的用户，升级到v0.5.0时需要注意以下几点：

1. 检查随机数生成相关的代码，确保与新的PRNG语义兼容
2. 更新依赖的NumPy和SciPy版本至最低要求以上
3. 如果使用了被废弃或移除的API，需要按照文档指引进行迁移
4. Mac x86用户需要考虑迁移到ARM架构或等待社区支持恢复

JAX团队持续致力于提供高性能的数值计算解决方案，本次更新在保持核心功能稳定的同时，对底层实现进行了重要优化，为未来的性能提升和功能扩展奠定了基础。