Index

_

_apply_weight_decay_inplace() (emerging_optimizers.mixin.WeightDecayMixin method)

A

AdaptiveMuon (class in emerging_optimizers.orthogonalized_optimizers)

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.update_functions) calculate_ademamix_update() (in module emerging_optimizers.scalar_optimizers.update_functions) calculate_laprop_update() (in module emerging_optimizers.scalar_optimizers.update_functions) calculate_lion_update() (in module emerging_optimizers.scalar_optimizers.update_functions) calculate_madam_update() (in module emerging_optimizers.scalar_optimizers.update_functions)

calculate_rmsprop_update() (in module emerging_optimizers.scalar_optimizers.update_functions) calculate_signum_update() (in module emerging_optimizers.scalar_optimizers.update_functions) calculate_sim_ademamix_update() (in module emerging_optimizers.scalar_optimizers.update_functions) conjugate() (in module emerging_optimizers.utils.eig) Conv1dFlatWeights (class in emerging_optimizers.utils.modules)

E

eigh_with_fallback() (in module emerging_optimizers.utils.eig) emerging_optimizers.legacy_soap.soap_utils module 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.isospectral module emerging_optimizers.riemannian_optimizers.normalized_optimizer module emerging_optimizers.scalar_optimizers module emerging_optimizers.scalar_optimizers.update_functions module emerging_optimizers.utils module emerging_optimizers.utils.eig module emerging_optimizers.utils.modules module extra_repr() (emerging_optimizers.utils.modules.Conv1dFlatWeights method)

F

forward() (emerging_optimizers.utils.modules.Conv1dFlatWeights method)

fp32_matmul_precision() (in module emerging_optimizers.utils)

G

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

get_eigenbasis_svd() (in module emerging_optimizers.legacy_soap.soap_utils) get_pg_rank() (in module emerging_optimizers.utils) get_pg_size() (in module emerging_optimizers.utils)

I

Iso (class in emerging_optimizers.riemannian_optimizers.isospectral)

L

LaProp (class in emerging_optimizers.scalar_optimizers)

left_polargrad_orth_fn() (in module emerging_optimizers.orthogonalized_optimizers) Lion (class in emerging_optimizers.scalar_optimizers)

M

module emerging_optimizers.legacy_soap.soap_utils 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.isospectral emerging_optimizers.riemannian_optimizers.normalized_optimizer emerging_optimizers.scalar_optimizers emerging_optimizers.scalar_optimizers.update_functions emerging_optimizers.utils emerging_optimizers.utils.eig emerging_optimizers.utils.modules

MOP (class in emerging_optimizers.orthogonalized_optimizers) Muon (class in emerging_optimizers.orthogonalized_optimizers) MuonHyperball (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) one_sided_polargrad_orth_fn() (in module emerging_optimizers.orthogonalized_optimizers)

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) PolarGrad (class in emerging_optimizers.orthogonalized_optimizers) post_step_inplace() (emerging_optimizers.scalar_optimizers.LaProp method) post_weight_update_fn_inplace() (emerging_optimizers.orthogonalized_optimizers.MuonHyperball method) (emerging_optimizers.orthogonalized_optimizers.OrthogonalizedOptimizer method)

pre_step_inplace() (emerging_optimizers.scalar_optimizers.LaProp method) pre_weight_update_fn_inplace() (emerging_optimizers.orthogonalized_optimizers.MuonHyperball method) (emerging_optimizers.orthogonalized_optimizers.OrthogonalizedOptimizer method) procrustes_step() (in module emerging_optimizers.psgd.procrustes_step) project_in() (in module emerging_optimizers.legacy_soap.soap) project_out() (in module emerging_optimizers.legacy_soap.soap)

R

REKLS (class in emerging_optimizers.legacy_soap.rekls)

right_polargrad_orth_fn() (in module emerging_optimizers.orthogonalized_optimizers)

S

Scion (class in emerging_optimizers.orthogonalized_optimizers) Signum (class in emerging_optimizers.scalar_optimizers) SimplifiedAdEMAMix (class in emerging_optimizers.scalar_optimizers) SingleMomentumOptimizer (class in emerging_optimizers.scalar_optimizers) SOAP (class in emerging_optimizers.legacy_soap.soap)

step() (emerging_optimizers.legacy_soap.soap.SOAP method) (emerging_optimizers.orthogonalized_optimizers.AdaptiveMuon method) (emerging_optimizers.orthogonalized_optimizers.OrthogonalizedOptimizer method) (emerging_optimizers.riemannian_optimizers.isospectral.Iso method) (emerging_optimizers.riemannian_optimizers.normalized_optimizer.ObliqueAdam method) (emerging_optimizers.riemannian_optimizers.normalized_optimizer.ObliqueSGD method)

T

TwoMomentsOptimizer (class in emerging_optimizers.scalar_optimizers)

U

uniformize_q_in_place() (in module emerging_optimizers.psgd.psgd_utils) update_eigenbasis_and_exp_avgs() (in module emerging_optimizers.legacy_soap.soap)

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

W

WeightDecayMixin (class in emerging_optimizers.mixin)