JAX v0.5.0发布:随机数生成器重大变更与多维FFT增强
JAX是由Google开发的高性能数值计算库,它结合了NumPy的易用性与自动微分、GPU/TPU加速等高级功能,特别适合机器学习研究和科学计算。JAX通过XLA编译器将Python/NumPy代码转换为高效的硬件加速指令,同时提供了自动微分和向量化等强大功能。
版本升级策略调整
JAX v0.5.0开始采用基于开发投入的版本控制方案(effort-based versioning)。这种版本控制方式更加注重反映开发团队的实际工作量和变更的重要性,而非严格遵循语义化版本控制。本次版本将"meso"版本号提升至v0.5.0,标志着包含了一项可能影响用户代码的破坏性变更。
重要破坏性变更
随机数生成器默认实现变更
JAX v0.5.0默认启用了jax_threefry_partitionable
随机数生成器实现。这一变更会影响PRNG(伪随机数生成器)密钥的语义,可能导致现有代码需要调整。开发团队已提供了详细的更新指南,帮助用户平滑过渡到新版本。
Mac x86架构支持终止
本次发布停止了对Mac x86平台wheel包的支持,主要基于两点考虑:
- Mac x86构建存在多个测试失败和崩溃问题,团队不愿发布不稳定的版本
- Mac x86硬件已停产,开发者难以获取设备进行问题修复
值得注意的是,Mac ARM架构仍然得到完全支持。社区成员若有意协助维护Mac x86平台,需要确保JAX测试套件在该平台上完全通过。
功能增强与改进
多维FFT扩展支持
JAX v0.5.0显著增强了快速傅里叶变换(FFT)功能:
fftn
、rfftn
、ifftn
和irfftn
函数现在支持超过3维的变换- 这一改进为高维信号处理和分析提供了更强大的工具
外部函数接口(FFI)增强
新增了jax.ffi.register_ffi_type_id
函数,允许用户:
- 在FFI中定义自定义状态
- 实现更灵活的与外部代码的交互
- 构建更复杂的跨语言集成方案
依赖要求更新
JAX v0.5.0调整了对核心科学计算库的最低版本要求:
- NumPy最低版本提升至1.25(支持至2025年6月)
- SciPy最低版本提升至1.11(支持至2025年6月)
API变更与优化
性能优化
jax.numpy.einsum
函数现在默认使用optimize='auto'
而非optimize='optimal'
,这一变更:
- 避免了在多参数情况下出现指数级增长的追踪时间
- 在大多数情况下仍能保持良好的性能
- 用户仍可显式指定
optimize='optimal'
以获得最佳性能
线性代数函数行为调整
jax.numpy.linalg.solve
不再支持右侧参数为批量1D数组的情况。需要此功能的用户应使用solve(a, b[..., None]).squeeze(-1)
来保持原有行为。
废弃与移除功能
已废弃功能
jax.interpreters.xla
模块中的abstractify
和pytype_aval_mappings
已移至jax.core
jax.scipy.special.lpmn
和jax.scipy.special.lpmn_values
函数已被废弃jax.extend.ffi
子模块已迁移至jax.ffi
已移除功能
jax_enable_memories
标志已被移除,其功能默认启用jax.lib.xla_client
中的Device
和XlaRuntimeError
已被移除jax.experimental.array_api
模块已被移除
开发者工具增强
AOT(提前编译)降低级的.as_text()
方法新增了debug_info
选项,可以:
- 在输出中包含调试信息
- 显示源代码位置等关键信息
- 帮助开发者更高效地诊断问题
升级建议
对于现有JAX用户,升级到v0.5.0时应当:
- 首先检查随机数生成相关的代码是否受影响
- 确认是否使用了已废弃或移除的API
- 在测试环境中验证关键功能是否正常
- 特别注意Mac x86用户需要考虑迁移到ARM架构
JAX持续演进的核心目标是提供更稳定、高效的数值计算平台,同时保持与科学计算生态系统的紧密集成。v0.5.0的发布标志着JAX在随机数生成和多维信号处理等关键领域的重要进步。
- DDeepSeek-V3.1-BaseDeepSeek-V3.1 是一款支持思考模式与非思考模式的混合模型Python00
- QQwen-Image-Edit基于200亿参数Qwen-Image构建,Qwen-Image-Edit实现精准文本渲染与图像编辑,融合语义与外观控制能力Jinja00
GitCode-文心大模型-智源研究院AI应用开发大赛
GitCode&文心大模型&智源研究院强强联合,发起的AI应用开发大赛;总奖池8W,单人最高可得价值3W奖励。快来参加吧~050CommonUtilLibrary
快速开发工具类收集,史上最全的开发工具类,欢迎Follow、Fork、StarJava04GitCode百大开源项目
GitCode百大计划旨在表彰GitCode平台上积极推动项目社区化,拥有广泛影响力的G-Star项目,入选项目不仅代表了GitCode开源生态的蓬勃发展,也反映了当下开源行业的发展趋势。06GOT-OCR-2.0-hf
阶跃星辰StepFun推出的GOT-OCR-2.0-hf是一款强大的多语言OCR开源模型,支持从普通文档到复杂场景的文字识别。它能精准处理表格、图表、数学公式、几何图形甚至乐谱等特殊内容,输出结果可通过第三方工具渲染成多种格式。模型支持1024×1024高分辨率输入,具备多页批量处理、动态分块识别和交互式区域选择等创新功能,用户可通过坐标或颜色指定识别区域。基于Apache 2.0协议开源,提供Hugging Face演示和完整代码,适用于学术研究到工业应用的广泛场景,为OCR领域带来突破性解决方案。00openHiTLS
旨在打造算法先进、性能卓越、高效敏捷、安全可靠的密码套件,通过轻量级、可剪裁的软件技术架构满足各行业不同场景的多样化要求,让密码技术应用更简单,同时探索后量子等先进算法创新实践,构建密码前沿技术底座!C0305- WWan2.2-S2V-14B【Wan2.2 全新发布|更强画质,更快生成】新一代视频生成模型 Wan2.2,创新采用MoE架构,实现电影级美学与复杂运动控制,支持720P高清文本/图像生成视频,消费级显卡即可流畅运行,性能达业界领先水平Python00
- GGLM-4.5-AirGLM-4.5 系列模型是专为智能体设计的基础模型。GLM-4.5拥有 3550 亿总参数量,其中 320 亿活跃参数;GLM-4.5-Air采用更紧凑的设计,拥有 1060 亿总参数量,其中 120 亿活跃参数。GLM-4.5模型统一了推理、编码和智能体能力,以满足智能体应用的复杂需求Jinja00
Yi-Coder
Yi Coder 编程模型,小而强大的编程助手HTML013
热门内容推荐
最新内容推荐
项目优选









