jax.numpy.put_along_axis

jax.numpy.put_along_axis(arr, indices, values, axis, inplace=True, *, mode=None)

通过匹配 1d 索引和数据切片将值放入目标数组。

numpy.put_along_axis() 的 JAX 实现。

numpy.put_along_axis() 的语义是原位修改数组，这对于 JAX 的不可变数组是不可能的。JAX 版本返回输入的修改副本，并添加 inplace 参数，用户必须将其设置为 False` 以提醒此 API 差异。

参数:

• arr (Array | ndarray | bool | number | bool | int | float | complex) – 将要放入值的数组。

• indices (Array | ndarray | bool | number | bool | int | float | complex) – 放置值的索引数组。

• values (Array | ndarray | bool | number | bool | int | float | complex) – 要放入数组的值数组。

• axis (int | None) – 放置值的轴。如果未指定，则将在应用索引之前展平数组。

• inplace (bool) – 必须设置为 False 以指示输入未原位修改，而是返回修改后的副本。

• mode (str | None) – 越界索引模式。有关 mode 选项的更多讨论，请参阅 jax.numpy.ndarray.at。

返回:

更新了指定条目的 a 的副本。

返回类型:

Array

另请参阅

• jax.numpy.put()：将元素放入给定索引的数组中。

• jax.numpy.place()：通过布尔掩码将元素放入数组中。

• jax.numpy.ndarray.at()：使用 NumPy 样式索引的数组更新。

• jax.numpy.take()：在给定索引处从数组中提取值。

• jax.numpy.take_along_axis()：沿轴从数组中提取值。

示例

>>> from jax import numpy as jnp
>>> a = jnp.array([[10, 30, 20], [60, 40, 50]])
>>> i = jnp.argmax(a, axis=1, keepdims=True)
>>> print(i)
[[1]
 [0]]
>>> b = jnp.put_along_axis(a, i, 99, axis=1, inplace=False)
>>> print(b)
[[10 99 20]
 [99 40 50]]