JAX

JAX

5|0 评价|0 收藏
访问网站
产品介绍
JAX 是一个用于高性能数组计算和程序转换的 Python 库，专为数值计算和机器学习而设计。
收录时间
2026年7月22日
免费网站AI 框架AI SDKAI 计算机视觉项目AI 大模型 APIAI 机器学习平台AI 模型部署AI 模型微调AI 模型管理AI 开源大模型AI 工具库
JAX 产品预览
访问网站

什么是 JAX？

JAX 是一个用于面向加速器的数组计算和程序转换的 Python 库，专为高性能数值计算和大规模机器学习而设计。它提供了熟悉的 NumPy 风格 API，并包含可组合的函数转换，用于编译、批处理、自动微分和并行化。相同的代码可以在多种后端上运行，包括 CPU、GPU 和 TPU。

如何使用 JAX？

  • 通过 pip 或 conda 安装 JAX，然后导入 jax 和 jax.numpy。使用 jax.jit 进行即时编译，jax.vmap 进行自动向量化，jax.grad 进行自动微分，jax.pmap 进行并行化。有关教程和 API 参考，请参阅官方文档。

JAX 的核心功能

  • 熟悉的 NumPy 风格 API
  • 可组合的函数转换（jit、vmap、grad、pmap）
  • 自动微分（前向和反向模式）
  • 即时编译
  • 自动向量化
  • 跨多个设备的并行化
  • 支持 CPU、GPU 和 TPU 后端

用户评价

5.0

/ 5

0 条评价

5 星
0
4 星
0
3 星
0
2 星
0
1 星
0

暂无评价，来写第一条吧

相似产品

JAX 推广组件

使用网站徽章展示您的产品已收录于 Hootool，帮助更多人按问题找到合适的 AI 工具。推广组件可轻松添加到主页或页脚。

Featured on

Hootool.ai