optax.power_iteration

optax.power_iteration#

optax.power_iteration(matrix: jax.typing.ArrayLike | Callable[[base.ArrayTree], base.ArrayTree], *, v0: base.ArrayTree | None = None, num_iters: jax.typing.ArrayLike = 100, error_tolerance: jax.typing.ArrayLike = 1e-06, precision: lax.Precision = Precision.HIGHEST, key: base.PRNGKey | None = None) → tuple[jax.typing.ArrayLike, base.ArrayTree][source]#

Power iteration algorithm.

This algorithm computes the dominant eigenvalue (i.e. the spectral radius) and its associated eigenvector of a diagonalizable matrix. This matrix can be given as an array or as a callable implementing a matrix-vector product.

Parameters:

  • matrix – a square matrix, either as an array or a callable implementing a matrix-vector product.

  • v0 – initial vector approximating the dominiant eigenvector. If matrix is an array of size (n, n), v0 must be a vector of size (n,). If instead matrix is a callable, then v0 must be a tree with the same structure as the input of this callable. If this argument is None and matrix is an array, then a random vector sampled from a uniform distribution in [-1, 1] is used as initial vector.

  • num_iters – Number of power iterations.

  • error_tolerance – Iterative exit condition. The procedure stops when the relative error of the estimate of the dominant eigenvalue is below this threshold.

  • precision – precision XLA related flag, the available options are: a) lax.Precision.DEFAULT (better step time, but not precise); b) lax.Precision.HIGH (increased precision, slower); c) lax.Precision.HIGHEST (best possible precision, slowest).

  • key – random key for the initialization of v0 when not given explicitly. When this argument is None, jax.random.PRNGKey(0) is used.

Returns:

A pair (eigenvalue, eigenvector), where eigenvalue is the dominant eigenvalue of matrix and eigenvector is its associated eigenvector.

References

Wikipedia contributors. Power iteration.

Note

If the matrix is not diagonalizable or the dominant eigenvalue is not unique, the algorithm may not converge.

Changed in version 0.2.2: matrix can be a callable. Reversed the order of the return parameters, from (eigenvector, eigenvalue) to (eigenvalue, eigenvector).