jax.numpy.fill_diagonal

jax.numpy.fill_diagonal(a, val, wrap=False, *, inplace=True)

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

JAX 实现的 numpy.fill_diagonal()。

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

参数:

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

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

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

• wrap (bool)

返回:

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

返回类型:

数组

示例

>>> 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)