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

入門指南

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

更多指南/資源

  • 使用者指南
    • 如何以 JAX 的思維思考
    • 效能分析計算
    • 效能分析裝置記憶體
    • 偵錯執行階段值
      • 編譯後的列印和中斷點
      • checkify 轉換
      • JAX 偵錯旗標
    • GPU 效能提示
    • 持續編譯快取
    • Pytrees
    • 錯誤
    • 預先降低階層和編譯
    • 匯出和序列化
      • 匯出和序列化階段性計算
      • 形狀多型
      • 與 TensorFlow 互通
    • 型別提升語意
    • Transfer guard
    • Pallas：JAX 核心語言
      • Pallas 快速入門
      • Grids 和 BlockSpecs
      • Pallas TPU
        • 使用 Pallas 撰寫 TPU 核心
        • 管線化
        • 矩陣乘法
        • 純量預取和區塊稀疏計算
        • 用於 TPU 的 Pallas 分散式運算
      • Pallas 設計筆記
        • Pallas 設計
        • Pallas 非同步操作
      • Pallas 變更記錄
    • 外部函式介面 (FFI)
    • 訓練簡單神經網路，使用 tensorflow/datasets 資料載入
    • 訓練簡單神經網路，使用 PyTorch 資料載入
    • 用於貝氏推論的自動批次處理
  • 進階指南
    • 分散式陣列和自動平行化
    • 使用 shard_map 的手動平行化
    • 多主機和多進程環境
    • 分散式資料載入
    • 自動微分食譜
    • 自訂導數規則
    • 使用 jax.checkpoint (又名 jax.remat) 控制自動微分的儲存值
    • JAX 中的廣義卷積
    • XLA 編譯器旗標列表
  • 開發者筆記
    • 貢獻 JAX
    • 從原始碼建置
    • 調查迴歸
    • Autodidax：從頭開始的 JAX 核心
    • JAX 增強提案 (JEP)
      • 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 Wrappers 的範圍
      • 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 提升警告
  • 公開 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.lax.index_take

目錄

  • index_take()

jax.lax.index_take#

jax.lax.index_take(src, idxs, axes)[原始碼]#
參數:

  • src (Array)

  • idxs (Array)

  • axes (Sequence[int])

回傳型別:

Array

上一頁

jax.lax.index_in_dim

下一頁

jax.lax.integer_pow

目錄
  • index_take()

由 JAX 作者群

© Copyright 2024, The JAX Authors.