jax.numpy.place

jax.numpy.place(arr, mask, vals, *, inplace=True)

根据掩码更新数组元素。

JAX 实现 numpy.place()。

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

参数:

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

• mask (Array | ndarray | bool | number | bool | int | float | complex | TypedNdArray) – 布尔掩码，大小与 arr 相同。

• vals (Array | ndarray | bool | number | bool | int | float | complex | TypedNdArray) – 要插入到 arr 中由掩码指示的位置的值。如果提供的数值过多，则会被截断。如果提供的数值不足，则会重复使用。

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

返回:

一个 arr 的副本，其中掩码值被设置为 vals 中的条目。

返回类型:

Array

另请参阅

• jax.numpy.put()：在数值索引处将元素放入数组。

• jax.numpy.ndarray.at()：使用 NumPy 风格的索引进行数组更新

示例

>>> x = jnp.zeros((3, 5), dtype=int)
>>> mask = (jnp.arange(x.size) % 3 == 0).reshape(x.shape)
>>> mask
Array([[ True, False, False,  True, False],
       [False,  True, False, False,  True],
       [False, False,  True, False, False]], dtype=bool)

放置标量值

>>> jnp.place(x, mask, 1, inplace=False)
Array([[1, 0, 0, 1, 0],
       [0, 1, 0, 0, 1],
       [0, 0, 1, 0, 0]], dtype=int32)

在这种情况下，jnp.place 类似于掩码数组更新语法

>>> x.at[mask].set(1)
Array([[1, 0, 0, 1, 0],
       [0, 1, 0, 0, 1],
       [0, 0, 1, 0, 0]], dtype=int32)

place 在放置数组值时有所不同。数组会被重复以填充掩码条目

>>> vals = jnp.array([1, 3, 5])
>>> jnp.place(x, mask, vals, inplace=False)
Array([[1, 0, 0, 3, 0],
       [0, 5, 0, 0, 1],
       [0, 0, 3, 0, 0]], dtype=int32)