jax.numpy.linalg.trace

jax.numpy.linalg.trace(x, /, *, offset=0, dtype=None)

计算矩阵的迹。

JAX 对 numpy.linalg.trace() 的实现。

参数:

• x (ArrayLike) – 形状为 (..., M, N) 的数组，其最内层的两个维度形成 MxN 的矩阵，用于计算迹。

• offset (int) – 主对角线的正或负偏移量（默认值：0）。

• dtype (DTypeLike | None | None) – 返回数组的数据类型（默认值：None）。如果为 None，则输出的 dtype 将与 x 的 dtype 匹配，对于整数类型，则会提升到默认精度。

返回值:

形状为 x.shape[:-2] 的批量迹的数组

返回类型:

数组

另请参阅

• jax.numpy.trace()：在 jax.numpy 命名空间中类似的 API。

示例

单个矩阵的迹

>>> x = jnp.array([[1,  2,  3,  4],
...                [5,  6,  7,  8],
...                [9, 10, 11, 12]])
>>> jnp.linalg.trace(x)
Array(18, dtype=int32)
>>> jnp.linalg.trace(x, offset=1)
Array(21, dtype=int32)
>>> jnp.linalg.trace(x, offset=-1, dtype="float32")
Array(15., dtype=float32)

批量迹

>>> x = jnp.arange(24).reshape(2, 3, 4)
>>> jnp.linalg.trace(x)
Array([15, 51], dtype=int32)