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在随机数生成和多维信号处理等关键领域的重要进步。
HunyuanImage-3.0
HunyuanImage-3.0 统一多模态理解与生成,基于自回归框架,实现文本生成图像,性能媲美或超越领先闭源模型00ops-transformer
本项目是CANN提供的transformer类大模型算子库,实现网络在NPU上加速计算。C++036Hunyuan3D-Part
腾讯混元3D-Part00GitCode-文心大模型-智源研究院AI应用开发大赛
GitCode&文心大模型&智源研究院强强联合,发起的AI应用开发大赛;总奖池8W,单人最高可得价值3W奖励。快来参加吧~0283Hunyuan3D-Omni
腾讯混元3D-Omni:3D版ControlNet突破多模态控制,实现高精度3D资产生成00Spark-Chemistry-X1-13B
科大讯飞星火化学-X1-13B (iFLYTEK Spark Chemistry-X1-13B) 是一款专为化学领域优化的大语言模型。它由星火-X1 (Spark-X1) 基础模型微调而来,在化学知识问答、分子性质预测、化学名称转换和科学推理方面展现出强大的能力,同时保持了强大的通用语言理解与生成能力。Python00GOT-OCR-2.0-hf
阶跃星辰StepFun推出的GOT-OCR-2.0-hf是一款强大的多语言OCR开源模型,支持从普通文档到复杂场景的文字识别。它能精准处理表格、图表、数学公式、几何图形甚至乐谱等特殊内容,输出结果可通过第三方工具渲染成多种格式。模型支持1024×1024高分辨率输入,具备多页批量处理、动态分块识别和交互式区域选择等创新功能,用户可通过坐标或颜色指定识别区域。基于Apache 2.0协议开源,提供Hugging Face演示和完整代码,适用于学术研究到工业应用的广泛场景,为OCR领域带来突破性解决方案。00- HHowToCook程序员在家做饭方法指南。Programmer's guide about how to cook at home (Chinese only).Dockerfile09
- PpathwayPathway is an open framework for high-throughput and low-latency real-time data processing.Python00
最新内容推荐
项目优选









