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

開始使用

  • 安裝
  • 快速入門
  • 教學
    • 快速入門
    • 主要概念
    • 即時編譯
    • 自動向量化
    • 自動微分
    • 偵錯簡介
    • 虛擬隨機數
    • 使用 pytrees
    • 平行程式設計簡介
    • 具狀態的計算
    • 使用 JIT 的控制流程和邏輯運算子
    • 進階自動微分
    • 外部回呼
    • 使用 jax.checkpoint (jax.remat) 的梯度檢查點
    • JAX 內部機制：primitives
    • JAX 內部機制：jaxpr 語言
  • 🔪 JAX - 尖銳之處 🔪
  • 常見問題 (FAQ)

更多指南/資源

  • 使用者指南
    • 如何以 JAX 思考
    • 分析計算效能
    • 分析裝置記憶體
    • 偵錯執行階段值
      • 編譯後的列印和中斷點
      • checkify 轉換
      • JAX 偵錯旗標
    • GPU 效能提示
    • 持久編譯快取
    • Pytrees
    • 錯誤
    • 預先降低和編譯
    • 匯出和序列化
      • 匯出和序列化 staged-out 計算
      • 形狀多型性
      • 與 TensorFlow 互通
    • 類型提升語意
    • Transfer guard
    • Pallas：JAX 核心語言
      • Pallas 快速入門
      • Grids 和 BlockSpecs
      • Pallas TPU
        • 使用 Pallas 撰寫 TPU 核心
        • 管線化
        • 矩陣乘法
        • 純量預取和區塊稀疏計算
        • Pallas 中用於 TPU 的分散式運算
      • Pallas 設計筆記
        • Pallas 設計
        • Pallas 非同步操作
      • Pallas 變更日誌
    • 外部函數介面 (FFI)
    • 訓練簡單神經網路，使用 tensorflow/datasets 資料載入
    • 訓練簡單神經網路，使用 PyTorch 資料載入
    • 貝氏推論的自動批次處理
  • 進階指南
    • 分散式陣列和自動平行化
    • 使用 shard_map 的手動平行化
    • 多主機和多進程環境
    • 分散式資料載入
    • 自動微分食譜
    • 自訂導數規則
    • 使用 jax.checkpoint (又名 jax.remat) 控制自動微分的已儲存值
    • JAX 中的廣義卷積
    • XLA 編譯器旗標清單
  • 開發人員筆記
    • 貢獻 JAX
    • 從原始碼建置
    • 調查迴歸
    • Autodidax：從頭開始的 JAX 核心
    • JAX 增強提案 (JEPs)
      • 263: JAX PRNG 設計
      • 2026: JAX 可轉換函數的自訂 JVP/VJP 規則
      • 4008: 自訂 VJP 和 `nondiff_argnums` 更新
      • 4410: Omnistaging
      • 9263: 類型化金鑰 & 可插拔 RNG
      • 9407: JAX 類型提升語意設計
      • 9419: Jax 和 Jaxlib 版本控制
      • 10657: JAX 中的序列副作用
      • 11830: `jax.remat` / `jax.checkpoint` 新實作
      • 12049: JAX 的類型註釋路線圖
      • 14273: 用於簡單的每裝置程式碼的 `shard_map` (`shmap`)
      • 15856: `jax.extend`，擴充模組
      • 17111: `shard_map`（和其他映射）的有效轉置
      • 18137: JAX NumPy & SciPy 包裝函式的範圍
      • 25516: 基於努力的版本控制
  • 擴充指南
    • 在 JAX 中撰寫自訂 Jaxpr 直譯器
    • 使用 C++ 和 CUDA 的 GPU 自訂操作
    • jax.extend 模組
      • jax.extend.core 模組
      • jax.extend.linear_util 模組
      • jax.extend.mlir 模組
      • jax.extend.random 模組
    • 以 JAX 為基礎建置
  • 筆記
    • API 相容性
    • Python 和 NumPy 版本支援政策
    • jax.Array 遷移
    • 非同步分派
    • 並行
    • GPU 記憶體配置
    • Rank promotion 警告
  • 公開 API：jax 套件
    • jax.numpy 模組
      • 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 模組
      • jax.scipy.stats.bernoulli.logpmf
      • jax.scipy.stats.bernoulli.pmf
      • jax.scipy.stats.bernoulli.cdf
      • jax.scipy.stats.bernoulli.ppf
    • jax.lax 模組
    • jax.random 模組
    • jax.sharding 模組
    • jax.debug 模組
    • jax.dlpack 模組
    • jax.distributed 模組
    • jax.dtypes 模組
    • jax.ffi 模組
    • jax.flatten_util 模組
    • jax.image 模組
    • jax.nn 模組
      • jax.nn.initializers 模組
    • jax.ops 模組
    • jax.profiler 模組
    • jax.stages 模組
    • jax.tree 模組
    • jax.tree_util 模組
    • jax.typing 模組
    • jax.export 模組
    • jax.extend 模組
      • jax.extend.core 模組
      • jax.extend.linear_util 模組
      • jax.extend.mlir 模組
      • jax.extend.random 模組
    • jax.example_libraries 模組
      • jax.example_libraries.optimizers 模組
      • jax.example_libraries.stax 模組
    • jax.experimental 模組
      • jax.experimental.checkify 模組
      • jax.experimental.compilation_cache 模組
      • jax.experimental.custom_partitioning 模組
      • jax.experimental.jet 模組
      • jax.experimental.key_reuse 模組
      • jax.experimental.mesh_utils 模組
      • jax.experimental.multihost_utils 模組
      • jax.experimental.pallas 模組
        • jax.experimental.pallas.mosaic_gpu 模組
        • jax.experimental.pallas.triton 模組
        • jax.experimental.pallas.tpu 模組
      • jax.experimental.pjit 模組
      • jax.experimental.serialize_executable 模組
      • jax.experimental.shard_map 模組
      • jax.experimental.sparse 模組
        • 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 模組
    • 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
  • 關於專案
  • 變更記錄
  • 詞彙表
  • .rst

jax.experimental.pallas.mosaic_gpu.ACC

目錄

  • ACC

jax.experimental.pallas.mosaic_gpu.ACC#

jax.experimental.pallas.mosaic_gpu.ACC[原始碼]#

別名為 WGMMAAccumulatorRef

上一頁

jax.experimental.pallas.mosaic_gpu.wgmma_wait

下一頁

jax.experimental.pallas.mosaic_gpu.GMEM

目錄
  • ACC

由 JAX 作者群

© 版權所有 2024，JAX 作者群。