jax.numpy.diagonal

内容

jax.numpy.diagonal#

jax.numpy.diagonal(a, offset=0, axis1=0, axis2=1)[source]#

返回数组的指定对角线。

JAX 实现 numpy.diagonal()。

JAX 版本始终返回输入的副本，尽管如果这在 JIT 编译中使用，编译器可能会避免复制。

参数：

  • a (ArrayLike) – 输入数组。必须至少是二维的。

  • offset (int) – 可选，默认为 0。主对角线的对角线偏移量。必须是静态整数值。可以是正数或负数。

  • axis1 (int) – 可选，默认为 0。获取对角线的第一个轴。

  • axis2 (int) –

    可选，默认值为 1。用于获取对角线的第二个轴。

    返回值

    对于二维输入，返回一个一维数组；对于 N 维输入，通常返回一个 N-1 维数组。

返回类型:

数组

另请参阅

示例

>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6],
...                [7, 8, 9]])
>>> jnp.diagonal(x)
Array([1, 5, 9], dtype=int32)
>>> jnp.diagonal(x, offset=1)
Array([2, 6], dtype=int32)
>>> jnp.diagonal(x, offset=-1)
Array([4, 8], dtype=int32)
内容