jax.numpy.fill_diagonal

jax.numpy.fill_diagonal#

jax.numpy.fill_diagonal(a, val, wrap=False, *, inplace=True)[源代码][源代码]#

返回一个数组的副本，其中对角线被覆盖。

JAX 实现的 numpy.fill_diagonal()。

The semantics of numpy.fill_diagonal() are to modify arrays in-place, which is not possible for JAX’s immutable arrays. The JAX version returns a modified copy of the input, and adds the inplace parameter which must be set to False` by the user as a reminder of this API difference.

参数:

  • a (ArrayLike) – 输入数组。必须满足 a.ndim >= 2。如果 a.ndim >= 3，则所有维度的大小必须相同。

  • val (ArrayLike) – 用于填充对角线的标量或数组。如果是数组，它将被展平并重复以填充对角线元素。

  • inplace (bool) – 必须设置为 False 以表明输入未就地修改，而是返回了一个修改后的副本。

  • wrap (bool)

返回:

a 的副本，其对角线设置为 val。

返回类型:

Array

示例

>>> x = jnp.zeros((3, 3), dtype=int)
>>> jnp.fill_diagonal(x, jnp.array([1, 2, 3]), inplace=False)
Array([[1, 0, 0],
       [0, 2, 0],
       [0, 0, 3]], dtype=int32)

与 numpy.fill_diagonal() 不同，输入 x 不会被修改。

如果对角线值包含过多条目，它将被截断

>>> jnp.fill_diagonal(x, jnp.arange(100, 200), inplace=False)
Array([[100,   0,   0],
       [  0, 101,   0],
       [  0,   0, 102]], dtype=int32)

如果对角线上的条目太少，它将被重复：

>>> x = jnp.zeros((4, 4), dtype=int)
>>> jnp.fill_diagonal(x, jnp.array([3, 4]), inplace=False)
Array([[3, 0, 0, 0],
       [0, 4, 0, 0],
       [0, 0, 3, 0],
       [0, 0, 0, 4]], dtype=int32)

对于非方阵数组，主方阵切片的对角线将被填充：

>>> x = jnp.zeros((3, 5), dtype=int)
>>> jnp.fill_diagonal(x, 1, inplace=False)
Array([[1, 0, 0, 0, 0],
       [0, 1, 0, 0, 0],
       [0, 0, 1, 0, 0]], dtype=int32)

而对于N维方阵，N维对角线将被填充：

>>> y = jnp.zeros((2, 2, 2))
>>> jnp.fill_diagonal(y, 1, inplace=False)
Array([[[1., 0.],
        [0., 0.]],

       [[0., 0.],
        [0., 1.]]], dtype=float32)
目录