Index

_

_apply_weight_decay_inplace() (emerging_optimizers.mixin.WeightDecayMixin method)

A

apply_kronecker_factors() (in module emerging_optimizers.psgd.psgd_kron_contractions)

apply_preconditioner() (in module emerging_optimizers.psgd.psgd_kron_contractions)

C

calculate_adam_update() (in module emerging_optimizers.scalar_optimizers) calculate_ademamix_update() (in module emerging_optimizers.scalar_optimizers) calculate_laprop_update() (in module emerging_optimizers.scalar_optimizers)

calculate_lion_update() (in module emerging_optimizers.scalar_optimizers) calculate_signum_update() (in module emerging_optimizers.scalar_optimizers) calculate_sim_ademamix_update() (in module emerging_optimizers.scalar_optimizers) conjugate() (in module emerging_optimizers.utils.eig)

E

eigh_with_fallback() (in module emerging_optimizers.utils.eig) emerging_optimizers.mixin module emerging_optimizers.orthogonalized_optimizers module emerging_optimizers.orthogonalized_optimizers.muon_utils module emerging_optimizers.psgd.procrustes_step module emerging_optimizers.psgd.psgd_kron_contractions module emerging_optimizers.psgd.psgd_utils module

emerging_optimizers.riemannian_optimizers.normalized_optimizer module emerging_optimizers.scalar_optimizers module emerging_optimizers.soap.soap module emerging_optimizers.soap.soap_utils module emerging_optimizers.utils module emerging_optimizers.utils.eig module

F

fp32_matmul_precision() (in module emerging_optimizers.utils)

G

get_eigenbasis_eigh() (in module emerging_optimizers.soap.soap_utils) get_eigenbasis_qr() (in module emerging_optimizers.soap.soap_utils)

get_pg_rank() (in module emerging_optimizers.utils) get_pg_size() (in module emerging_optimizers.utils)

I

init_kronecker_factors() (in module emerging_optimizers.soap.soap)

M

met_approx_eigvals_criteria() (in module emerging_optimizers.utils.eig) module emerging_optimizers.mixin emerging_optimizers.orthogonalized_optimizers emerging_optimizers.orthogonalized_optimizers.muon_utils emerging_optimizers.psgd.procrustes_step emerging_optimizers.psgd.psgd_kron_contractions emerging_optimizers.psgd.psgd_utils emerging_optimizers.riemannian_optimizers.normalized_optimizer emerging_optimizers.scalar_optimizers emerging_optimizers.soap.soap emerging_optimizers.soap.soap_utils emerging_optimizers.utils emerging_optimizers.utils.eig

Muon (class in emerging_optimizers.orthogonalized_optimizers)

N

newton_schulz() (in module emerging_optimizers.orthogonalized_optimizers.muon_utils) newton_schulz_step() (in module emerging_optimizers.orthogonalized_optimizers.muon_utils)

newton_schulz_tp() (in module emerging_optimizers.orthogonalized_optimizers.muon_utils) norm_lower_bound_skew() (in module emerging_optimizers.psgd.psgd_utils) norm_lower_bound_spd() (in module emerging_optimizers.psgd.psgd_utils)

O

ObliqueAdam (class in emerging_optimizers.riemannian_optimizers.normalized_optimizer) ObliqueSGD (class in emerging_optimizers.riemannian_optimizers.normalized_optimizer)

orthogonal_iteration() (in module emerging_optimizers.utils.eig) orthogonalize() (emerging_optimizers.orthogonalized_optimizers.OrthogonalizedOptimizer method) OrthogonalizedOptimizer (class in emerging_optimizers.orthogonalized_optimizers)

P

partial_contraction() (in module emerging_optimizers.psgd.psgd_kron_contractions)

precondition() (in module emerging_optimizers.soap.soap) procrustes_step() (in module emerging_optimizers.psgd.procrustes_step)

S

Scion (class in emerging_optimizers.orthogonalized_optimizers) SOAP (class in emerging_optimizers.soap.soap) step() (emerging_optimizers.orthogonalized_optimizers.OrthogonalizedOptimizer method) (emerging_optimizers.riemannian_optimizers.normalized_optimizer.ObliqueAdam method) (emerging_optimizers.riemannian_optimizers.normalized_optimizer.ObliqueSGD method) (emerging_optimizers.soap.soap.SOAP method)

U

uniformize_q_in_place() (in module emerging_optimizers.psgd.psgd_utils) update_eigenbasis_and_momentum() (in module emerging_optimizers.soap.soap)

update_kronecker_factors() (in module emerging_optimizers.soap.soap) update_kronecker_factors_kl_shampoo() (in module emerging_optimizers.soap.soap)

W

WeightDecayMixin (class in emerging_optimizers.mixin)