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