jax.numpy.linalg.matrix_power

jax.numpy.linalg.matrix_power#

jax.numpy.linalg.matrix_power(a, n)[source]#

将方阵提升到整数幂。

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

参数：

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

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

返回：

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

返回类型：

Array

示例

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

负幂等价于逆矩阵的乘积

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