jax.numpy.trace

jax.numpy.trace(a, offset=0, axis1=0, axis2=1, dtype=None, out=None)

计算沿给定轴的输入对角线之和。

JAX 实现的 numpy.trace()。

参数:

• a (ArrayLike) – 输入数组。必须具有 a.ndim >= 2。

• offset (int | ArrayLike) – 可选，整数，默认值=0。主对角线的对角线偏移。可以是正数或负数。

• axis1 (int) – 可选，默认值=0。沿其取对角线和的第一个轴。必须是静态整数值。

• axis2 (int) – 可选，默认值=1。沿其取对角线和的第二个轴。必须是静态整数值。

• dtype (DTypeLike | None) – 可选。输出数组的 dtype。应在 JIT 编译中作为静态参数提供。

• out (None) – JAX 未使用。

返回值:

一个维度为 x.ndim-2 的数组，包含沿轴 (axis1, axis2) 的对角线元素之和

返回类型:

Array

另请参阅

• jax.numpy.diag()：返回指定的对角线或构造对角线数组

• jax.numpy.diagonal()：返回数组的指定对角线。

• jax.numpy.diagflat()：返回一个 2-D 数组，其扁平化的输入数组布局在对角线上。

示例

>>> x = jnp.arange(1, 9).reshape(2, 2, 2)
>>> x
Array([[[1, 2],
        [3, 4]],

       [[5, 6],
        [7, 8]]], dtype=int32)
>>> jnp.trace(x)
Array([ 8, 10], dtype=int32)
>>> jnp.trace(x, offset=1)
Array([3, 4], dtype=int32)
>>> jnp.trace(x, axis1=1, axis2=2)
Array([ 5, 13], dtype=int32)
>>> jnp.trace(x, offset=1, axis1=1, axis2=2)
Array([2, 6], dtype=int32)