jax.numpy.linalg.matrix_power#

jax.numpy.linalg.matrix_power(a, n)[源代码]#

将一个方阵提升到整数幂。

numpy.linalg.matrix_power()的 JAX 实现，通过重复平方实现。

参数:

  • a (ArrayLike) – 形状为 (..., M, M) 的数组，将提升到 n 幂。

  • n (int) – 矩阵应该提升到的整数指数。

返回:

形状为 (..., M, M) 的数组，包含 a 的 n 次方矩阵。

返回类型:

数组

示例

>>> a = jnp.array([[1., 2.],
...                [3., 4.]])
>>> jnp.linalg.matrix_power(a, 3)
Array([[ 37.,  54.],
       [ 81., 118.]], dtype=float32)
>>> a @ a @ a  # equivalent evaluated directly
Array([[ 37.,  54.],
       [ 81., 118.]], dtype=float32)

这也支持零次幂

>>> jnp.linalg.matrix_power(a, 0)
Array([[1., 0.],
       [0., 1.]], dtype=float32)

并且也支持负次幂

>>> with jnp.printoptions(precision=3):
...   jnp.linalg.matrix_power(a, -2)
Array([[ 5.5 , -2.5 ],
       [-3.75,  1.75]], dtype=float32)

负次幂等同于逆矩阵的 matmul

>>> inv_a = jnp.linalg.inv(a)
>>> with jnp.printoptions(precision=3):
...   inv_a @ inv_a
Array([[ 5.5 , -2.5 ],
       [-3.75,  1.75]], dtype=float32)