Skip to content

← JAX Numerical Computing step 5 of 25

Medium Primitives

linalg.matrix_power

Jump to the task

Repeat a linear transformation

If A @ x advances a state by one step, then A @ A @ x advances it by two steps. The matrix A @ A is written mathematically as A squared. Use jnp.linalg.matrix_power(A, n) to compute integer matrix powers.

For a positive integer n, the result is a product of n copies of A. This differs from A ** n, which raises each entry separately to that power. A matrix power combines rows and columns through multiplication.

Worked example

import jax.numpy as jnp

A = jnp.array([[0.0, 1.0],
               [1.0, 0.0]])
twice = jnp.linalg.matrix_power(A, 2)
# [[1.0, 0.0], [0.0, 1.0]]
three_times = jnp.linalg.matrix_power(A, 3)
# [[0.0, 1.0], [1.0, 0.0]]
elementwise = A ** 2
# [[0.0, 1.0], [1.0, 0.0]]

This matrix swaps two coordinates. Two swaps restore the original vector, so its square is the identity matrix. Three swaps have the same effect as one. Squaring the entries does not capture that behavior.

Zero and negative powers

The zeroth matrix power is the identity, including when A is singular. It represents applying the transformation zero times. For a negative exponent, A must be invertible: power -1 is its inverse, and power -2 is the inverse multiplied by itself.

For example, diag([2, 4]) raised to -2 is diag([0.25, 0.0625]). A singular matrix cannot support negative powers; numerical routines may produce NaNs or infinities rather than a useful inverse.

Problem

Implement matrix_power(A, n) using jnp.linalg.matrix_power.

  • A is a square float32 JAX array of shape (m, m).
  • n is a whole-number scalar supplied as a float. It can be positive, zero, or negative. Convert it to a Python integer for the eager call.
  • When n is negative, the input matrix is nonsingular.
  • Return a matrix of shape (m, m) representing the requested power.

Do not round a fractional exponent or compute elementwise powers. Fractional matrix powers are a different operation and are not requested.

Exponents under JIT

This API requires the exponent to be known when it is traced. The reason is the computation chosen by the integer exponent, not a changing output shape: every power still has shape (m, m).

int(n) works for the concrete scalar supplied to the eager function. It cannot extract an integer from a dynamic tracer. For a compiled adaptation, pass a Python exponent as static, for example with jax.jit(matrix_power, static_argnames=("n",)). Different static exponents can require separate compilations.

Repeated squaring reduces the number of matrix multiplications for large positive exponents. It does not prevent rounding error or overflow; negative powers can also be sensitive to poor conditioning.