emerging_optimizers.utils#

emerging_optimizers.utils.fp32_matmul_precision(
precision: str = 'highest',
) → Generator[None, None, None][source]#

Context manager for setting the precision of matmuls.

Parameters:

precision – Precision of matmuls (defaults to “highest”)

emerging_optimizers.utils.get_pg_rank(
group: ProcessGroup | None = None,
) → int[source]#

Get rank for a distributed group with fallback

emerging_optimizers.utils.get_pg_size(
group: ProcessGroup | None = None,
) → int[source]#

Get world size for a distributed group with fallback

emerging_optimizers.utils.eig#

emerging_optimizers.utils.eig.adaptive_early_exit_criteria(
approx_eigenvalues_matrix: Tensor,
tolerance: float,
) → bool[source]#

Evaluates if a criteria using approximated eigenvalues is below or equal to the tolerance.

approx_eigenvalues_matrix is a matrix created from the approximated eigenvectors and the symmetric matrix that is being eigendecomposed. We check if the ratio of the diagonal norm to the matrix norm is greater than or equal to (1 - tolerance).

Parameters:

  • approx_eigenvalues_matrix – The symmetric matrix whose eigenvalues is being eigendecomposed.

  • tolerance – The tolerance for the early exit criteria, the min relative error between diagonal norm and matrix norm of the approximated eigenvalues and the diagonal.

Returns:

True if the criteria is below or equal to the tolerance, False otherwise.

Return type:

bool

emerging_optimizers.utils.eig.eig_orthogonal_iteration(
x: Tensor,
approx_eigenvectors: Tensor,
max_iterations: int = 1,
tolerance: float = 0.01,
) → tuple[Tensor, Tensor][source]#

Approximately compute the eigendecomposition of a symmetric matrix by performing the orthogonal iteration algorithm.

Orthogonal or subspace iteration uses iterative power iteration and QR decomposition to update the approximated eigenvectors. When the initial estimate is the zero matrix, the eigendecomposition is computed using eigh_with_fallback.

Based on Purifying Shampoo (https://www.arxiv.org/abs/2506.03595), we use an early exit criteria to stop the QR iterations. This generalizes SOAP’s algorithm of 1 step of power iteration for updating the eigenbasis.

Parameters:

  • x – tensor of shape (n, n) where x is a symmetric or Hermitian matrix.

  • approx_eigenvectors – The current estimate of the eigenvectors of x. If None or a zero matrix, falls back to using eigh_with_fallback.

  • max_iterations – The maximum number of iterations to perform. (Default: 1)

  • tolerance – The tolerance for determining convergence in terms of the norm of the off-diagonal elements of the approximated eigenvalues. (Default: 0.01)

Returns:

A tuple containing the approximated eigenvalues and eigenvectors matrix of the input matrix A.

Return type:

tuple[Tensor, Tensor]

emerging_optimizers.utils.eig.eigh_with_fallback(
x: Tensor,
force_double: bool = False,
eps: float | None = None,
output_dtype: dtype | None = None,
) → tuple[Tensor, Tensor][source]#

torch.linalg.eigh() function with double precision fallback

Unified wrapper over eigh() function with automatic fallback and force double precision options. Automatically falls back to double precision on failure and returns eigenvalues in descending order. Default 2nd argument of eigh UPLO is ‘L’.

Parameters:

  • x – Tensor of shape (, n, n) where “” is zero or more batch dimensions consisting of symmetric or Hermitian matrices.

  • force_double – Force double precision computation. Default False.

  • eps – Small offset for numerical stability. If None, uses dtype-appropriate values (1e-7 for float32, 1e-15 for float64). Default None.

  • output_dtype – Desired output dtype. If None, uses input dtype. Default None.

Returns:

Eigenvalues and eigenvectors tuple (eigenvalues in descending order).

Return type:

tuple[Tensor, Tensor]