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 – 要加载的引用。

• idx – 要使用的索引器。

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

• other – 用于掩码为 False 的索引的可选值。

• cache_modifier – 待文档化。

• eviction_policy – 待文档化。

• volatile – 待文档化。

返回类型:

jax.Array