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)[source]#

从给定索引加载数组。

如果既未指定 mask 也未指定 other，则此函数具有与 JAX 中的 x_ref_or_view[idx] 相同的语义。

参数:

  • x_ref_or_view – 要加载的引用。

  • idx – 要使用的索引器。

  • mask – 可选的布尔掩码，用于指定要加载的索引。如果 mask 为 False 且未给出 other，则无法对结果数组中的值进行任何假设。

  • other – 可选值，用于 mask 为 False 的索引。

  • cache_modifier – 待撰写文档。

  • eviction_policy – 待撰写文档。

  • volatile – 待撰写文档。

返回类型:

jax.Array