jax.numpy.put

内容

jax.numpy.put#

jax.numpy.put(a, ind, v, mode=None, *, inplace=True)[source]#

将元素放入给定索引处的数组中。

JAX 实现的 numpy.put().

numpy.put() 的语义是就地修改数组，这对于 JAX 的不可变数组是不可能的。JAX 版本返回一个修改后的输入副本，并添加了 inplace 参数，该参数必须由用户设置为 False` 作为此 API 差异的提醒。

参数:

  • a (ArrayLike) – 将放置值的数组。

  • ind (ArrayLike) – 将放置值的扁平数组上的索引数组。

  • v (ArrayLike) – 要放入数组的值数组。

  • mode (str | None | None) –

    指定如何处理越界索引的字符串。支持的值

    • "clip" （默认值）：将越界索引裁剪到最终索引。

    • "wrap": 将超出边界索引包裹到数组的开头。

  • inplace (bool) – 必须设置为 False 以指示输入不会被就地修改，而是返回一个修改后的副本。

返回值:

一个包含指定条目更新的 a 的副本。

返回值类型:

数组

参见

示例

>>> x = jnp.zeros(5, dtype=int)
>>> indices = jnp.array([0, 2, 4])
>>> values = jnp.array([10, 20, 30])
>>> jnp.put(x, indices, values, inplace=False)
Array([10,  0, 20,  0, 30], dtype=int32)

这等效于以下 jax.numpy.ndarray.at 索引语法

>>> x.at[indices].set(values)
Array([10,  0, 20,  0, 30], dtype=int32)

处理超出边界索引有两种模式。默认情况下它们会被剪切

>>> indices = jnp.array([0, 2, 6])
>>> jnp.put(x, indices, values, inplace=False, mode='clip')
Array([10,  0, 20,  0, 30], dtype=int32)

或者，它们可以被包裹到数组的开头

>>> jnp.put(x, indices, values, inplace=False, mode='wrap')
Array([10,  30, 20,  0, 0], dtype=int32)

对于 N 维输入，索引引用的是扁平数组

>>> x = jnp.zeros((3, 5), dtype=int)
>>> indices = jnp.array([0, 7, 14])
>>> jnp.put(x, indices, values, inplace=False)
Array([[10,  0,  0,  0,  0],
       [ 0,  0, 20,  0,  0],
       [ 0,  0,  0,  0, 30]], dtype=int32)
内容