jax.experimental.pallas.load#

jax.experimental.pallas.load(x_ref_or_view, idx, *, mask=None, other=None, cache_modifier=None, eviction_policy=None, volatile=False)[原始碼]#

傳回從給定索引載入的陣列。

如果未指定 mask 和 other，則此函數具有與 JAX 中的 x_ref_or_view[idx] 相同的語意。

參數:

  • x_ref_or_view – 要從中載入的 ref。

  • idx – 要使用的索引器。

  • mask – 一個可選的布林遮罩，用於指定要載入的索引。如果 mask 為 False 且未給定 other，則無法對結果陣列中的值做出任何假設。

  • other – 一個可選值，用於 mask 為 False 的索引。

  • cache_modifier – 待文件化。

  • eviction_policy – 待文件化。

  • volatile – 待文件化。

傳回型別:

jax.Array