跳到主要内容
Ctrl+K
JAX documentation - Home

入门指南

  • 安装
  • 快速入门
  • 教程
    • 快速入门
    • 关键概念
    • 即时编译
    • 自动向量化
    • 自动微分
    • 调试简介
    • 伪随机数
    • 使用 pytrees
    • 并行编程简介
    • 有状态计算
    • JIT 的控制流和逻辑运算符
    • 高级自动微分
    • 外部回调
    • 使用 jax.checkpoint (jax.remat) 的梯度检查点
    • JAX 内部原理：原语
    • JAX 内部原理：jaxpr 语言
  • 🔪 JAX - 尖锐部分 🔪
  • 常见问题 (FAQ)

更多指南/资源

  • 用户指南
    • 如何用 JAX 思考
    • 计算性能分析
    • 设备内存性能分析
    • 调试运行时值
      • 编译后的打印和断点
      • checkify 转换
      • JAX 调试标志
    • GPU 性能提示
    • 持久化编译缓存
    • Pytrees
    • 错误
    • 提前 lowering 和编译
    • 导出和序列化
      • 导出和序列化 staged-out 计算
      • 形状多态性
      • 与 TensorFlow 互操作
    • 传输保护
    • Pallas：一种 JAX 内核语言
      • Pallas 快速入门
      • 软件流水线
      • 网格和 BlockSpecs
      • Pallas TPU
        • 使用 Pallas 编写 TPU 内核
        • 流水线
        • 矩阵乘法
        • 标量预取和块稀疏计算
        • Pallas 中用于 TPU 的分布式计算
      • Pallas:Mosaic GPU
        • 使用 Pallas 编写 Mosaic GPU 内核
      • Pallas 设计笔记
        • Pallas 设计
        • Pallas 异步操作
      • Pallas 变更日志
    • 外部函数接口 (FFI)
    • 训练一个简单的神经网络，使用 tensorflow/datasets 数据加载
    • 训练一个简单的神经网络，使用 PyTorch 数据加载
    • 用于贝叶斯推断的自动批处理
  • 高级指南
    • 分布式数组和自动并行化
    • 显式分片（也称为“类型分片”）
    • 使用 shard_map 的手动并行
    • 多主机和多进程环境
    • 分布式数据加载
    • 自动微分 Cookbook
    • 自定义导数规则
    • 使用 jax.checkpoint (或 jax.remat) 控制自动微分的保存值
    • JAX 中的广义卷积
    • XLA 编译器标志列表
  • 开发者笔记
    • 为 JAX 做贡献
    • 从源码构建
    • 调查回归
    • Autodidax：从零开始的 JAX 核心
    • Autodidax2，第 1 部分：再次从零开始的 JAX
    • JAX 增强提案 (JEPs)
      • 263: JAX PRNG Design
      • 2026: Custom JVP/VJP rules for JAX-transformable functions
      • 4008: Custom VJP and `nondiff_argnums` update
      • 4410: Omnistaging
      • 9263: Typed keys & pluggable RNGs
      • 9407: Design of Type Promotion Semantics for JAX
      • 9419: Jax and Jaxlib versioning
      • 10657: Sequencing side-effects in JAX
      • 11830: `jax.remat` / `jax.checkpoint` new implementation
      • 12049: Type Annotation Roadmap for JAX
      • 14273: `shard_map` (`shmap`) for simple per-device code
      • 15856: `jax.extend`, an extensions module
      • 17111: Efficient transposition of `shard_map` (and other maps)
      • 18137: Scope of JAX NumPy & SciPy Wrappers
      • 25516: Effort-based versioning
  • 扩展指南
    • 在 JAX 中编写自定义 Jaxpr 解释器
    • jax.extend module
      • jax.extend.core module
      • jax.extend.linear_util module
      • jax.extend.mlir module
      • jax.extend.random module
    • 基于 JAX 构建
  • 笔记
    • API 兼容性
    • Python 和 NumPy 版本支持策略
    • 异步调度
    • 并发
    • GPU 内存分配
    • 秩提升警告
    • 类型提升语义
    • 默认 dtypes 和 X64 标志
  • Public API: jax package
    • jax.numpy module
      • jax.numpy.fft.fft
      • jax.numpy.fft.fft2
      • jax.numpy.fft.fftfreq
      • jax.numpy.fft.fftn
      • jax.numpy.fft.fftshift
      • jax.numpy.fft.hfft
      • jax.numpy.fft.ifft
      • jax.numpy.fft.ifft2
      • jax.numpy.fft.ifftn
      • jax.numpy.fft.ifftshift
      • jax.numpy.fft.ihfft
      • jax.numpy.fft.irfft
      • jax.numpy.fft.irfft2
      • jax.numpy.fft.irfftn
      • jax.numpy.fft.rfft
      • jax.numpy.fft.rfft2
      • jax.numpy.fft.rfftfreq
      • jax.numpy.fft.rfftn
    • jax.scipy module
      • jax.scipy.stats.bernoulli.logpmf
      • jax.scipy.stats.bernoulli.pmf
      • jax.scipy.stats.bernoulli.cdf
      • jax.scipy.stats.bernoulli.ppf
    • jax.lax module
    • jax.random module
    • jax.sharding module
    • jax.debug module
    • jax.dlpack module
    • jax.distributed module
    • jax.dtypes module
    • jax.ffi module
    • jax.flatten_util module
    • jax.image module
    • jax.nn module
      • jax.nn.initializers module
    • jax.ops module
    • jax.profiler module
    • jax.stages module
    • jax.test_util module
    • jax.tree module
    • jax.tree_util module
    • jax.typing module
    • jax.export module
    • jax.extend module
      • jax.extend.core module
      • jax.extend.linear_util module
      • jax.extend.mlir module
      • jax.extend.random module
    • jax.example_libraries module
      • jax.example_libraries.optimizers module
      • jax.example_libraries.stax module
    • jax.experimental module
      • jax.experimental.checkify module
      • jax.experimental.compilation_cache module
      • jax.experimental.custom_dce module
      • jax.experimental.custom_partitioning module
      • jax.experimental.jet module
      • jax.experimental.key_reuse module
      • jax.experimental.mesh_utils module
      • jax.experimental.multihost_utils module
      • jax.experimental.pallas module
        • jax.experimental.pallas.mosaic_gpu module
        • jax.experimental.pallas.triton module
        • jax.experimental.pallas.tpu module
      • jax.experimental.pjit module
      • jax.experimental.serialize_executable module
      • jax.experimental.shard_map module
      • jax.experimental.sparse module
        • jax.experimental.sparse.BCOO
        • jax.experimental.sparse.bcoo_broadcast_in_dim
        • jax.experimental.sparse.bcoo_concatenate
        • jax.experimental.sparse.bcoo_dot_general
        • jax.experimental.sparse.bcoo_dot_general_sampled
        • jax.experimental.sparse.bcoo_dynamic_slice
        • jax.experimental.sparse.bcoo_extract
        • jax.experimental.sparse.bcoo_fromdense
        • jax.experimental.sparse.bcoo_gather
        • jax.experimental.sparse.bcoo_multiply_dense
        • jax.experimental.sparse.bcoo_multiply_sparse
        • jax.experimental.sparse.bcoo_update_layout
        • jax.experimental.sparse.bcoo_reduce_sum
        • jax.experimental.sparse.bcoo_reshape
        • jax.experimental.sparse.bcoo_slice
        • jax.experimental.sparse.bcoo_sort_indices
        • jax.experimental.sparse.bcoo_squeeze
        • jax.experimental.sparse.bcoo_sum_duplicates
        • jax.experimental.sparse.bcoo_todense
        • jax.experimental.sparse.bcoo_transpose
    • jax.lib module
    • jax.Array.addressable_shards
    • jax.Array.all
    • jax.Array.any
    • jax.Array.argmax
    • jax.Array.argmin
    • jax.Array.argpartition
    • jax.Array.argsort
    • jax.Array.astype
    • jax.Array.at
    • jax.Array.choose
    • jax.Array.clip
    • jax.Array.compress
    • jax.Array.committed
    • jax.Array.conj
    • jax.Array.conjugate
    • jax.Array.copy
    • jax.Array.copy_to_host_async
    • jax.Array.cumprod
    • jax.Array.cumsum
    • jax.Array.device
    • jax.Array.diagonal
    • jax.Array.dot
    • jax.Array.dtype
    • jax.Array.flat
    • jax.Array.flatten
    • jax.Array.global_shards
    • jax.Array.imag
    • jax.Array.is_fully_addressable
    • jax.Array.is_fully_replicated
    • jax.Array.item
    • jax.Array.itemsize
    • jax.Array.max
    • jax.Array.mean
    • jax.Array.min
    • jax.Array.nbytes
    • jax.Array.ndim
    • jax.Array.nonzero
    • jax.Array.prod
    • jax.Array.ptp
    • jax.Array.ravel
    • jax.Array.real
    • jax.Array.repeat
    • jax.Array.reshape
    • jax.Array.round
    • jax.Array.searchsorted
    • jax.Array.shape
    • jax.Array.sharding
    • jax.Array.size
    • jax.Array.sort
    • jax.Array.squeeze
    • jax.Array.std
    • jax.Array.sum
    • jax.Array.swapaxes
    • jax.Array.take
    • jax.Array.to_device
    • jax.Array.trace
    • jax.Array.transpose
    • jax.Array.var
    • jax.Array.view
    • jax.Array.T
    • jax.Array.mT
  • 关于项目
  • 变更日志
  • 术语表
  • 配置选项
  • 用户指南
  • Pallas：一种 JAX 内核语言
  • Pallas 设计笔记
  • .rst

Pallas 设计笔记

Pallas 设计笔记#

设计

  • Pallas 设计
    • 简介
    • Pallas：为内核扩展 JAX
  • Pallas 异步操作
    • 背景 + 动机
    • 我们如何实现分解的 Pallas 异步操作 (在 TPU 上)？
    • 为什么 *这* 不起作用？
    • 状态的反击
    • 结论

上一页

使用 Pallas 编写 Mosaic GPU 内核

下一页

Pallas 设计

由 JAX 作者编写

© 版权所有 2024，JAX 作者。