heavyball 1.7.1__py3-none-any.whl → 2.0.0__py3-none-any.whl

heavyball-2.0.0.dist-info/METADATA

@@ -0,0 +1,122 @@
+Metadata-Version: 2.4
+Name: heavyball
+Version: 2.0.0
+Summary: Efficient Optimizers
+Author-email: HeavyBall Authors <github.heavyball@nestler.sh>
+Project-URL: source, https://github.com/HomebrewML/HeavyBall
+Project-URL: tracker, https://github.com/HomebrewML/HeavyBall/issues
+Keywords: torch,optimizer,muon,soap,psgd
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Natural Language :: English
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: opt-einsum>=3.4.0
+Requires-Dist: torch>=2.7.0
+Requires-Dist: numpy
+Provides-Extra: dev
+Requires-Dist: pre-commit; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: matplotlib; extra == "dev"
+Requires-Dist: seaborn; extra == "dev"
+Requires-Dist: hyperopt; extra == "dev"
+Requires-Dist: pandas; extra == "dev"
+Requires-Dist: typer; extra == "dev"
+Requires-Dist: optuna; extra == "dev"
+Requires-Dist: optunahub; extra == "dev"
+Requires-Dist: botorch; extra == "dev"
+Requires-Dist: hebo; extra == "dev"
+Dynamic: license-file
+
+# heavyball
+
+[![PyPI version](https://img.shields.io/pypi/v/heavyball?color=blue)][pypi]
+[![License](https://img.shields.io/badge/license-BSD--3--Clause-blue.svg)][license]
+
+_High-performance, extensible, chainable optimizers for PyTorch._
+
+## Why heavyball
+
+- **Lightning-Fast Training**: Batched `foreach` operations deliver significant speedups on large models.
+- **Adaptive & Extensible**: Built-in AdamW, RMSprop, Schedule-Free algorithms, and PaLM-inspired schedules.
+- **Plug-and-Play**: Drop-in replacements for `torch.optim` with seamless integration.
+- **Customizable**: Chainable API lets you compose optimizers and transforms (MARS correction, cautious updates, orthogonal updates).
+- **Battle-Tested**: Extensive benchmarks and real-world examples included.
+
+## Key Features
+
+- Foreach-based optimizers: `ForeachAdamW`, `ForeachRMSprop`, `ForeachSFAdamW`, `Muon`, `ADOPT`, `MSAM` (Momentum SAM), …
+- Schedule-Free optimizers with dynamic learning rate adaptation.
+- Advanced update rules: MARS correction, cautious updates, PaLM beta2 scheduling.
+- Chainable transforms for custom optimization recipes.
+- Comprehensive benchmark suite (`benchmark/`).
+- Detailed documentation and example-driven tutorials.
+
+## Quickstart
+
+**Install:**
+```bash
+pip install heavyball
+```
+
+**Basic usage:**
+```python
+import torch
+from torch import nn
+from heavyball import ForeachAdamW
+
+model = nn.Sequential(
+    nn.Linear(128, 64), nn.ReLU(), nn.Linear(64, 10)
+)
+optimizer = ForeachAdamW(model.parameters(), lr=1e-3)
+
+for data, target in dataloader:
+    optimizer.zero_grad()
+    output = model(data)
+    loss = torch.nn.functional.cross_entropy(output, target)
+    loss.backward()
+    optimizer.step()
+```
+
+## Benchmarks
+
+> Reproduce benchmarks with:
+> ```bash
+> python3 -m benchmark.run_all_benchmarks --opt ForeachSOAP --opt LaProp --opt AdamW --opt Muon --opt ForeachCachedNewtonPSGD  --opt RMSprop --opt OrthoLaProp --opt ForeachSFAdamW --opt ForeachADOPT --opt LaPropOrtho --opt CachedPSGDKron --opt SignLaProp --opt ForeachSOLP --opt PSGDLRA --opt NewtonPSGDLRA --opt NewtonHybrid2PSGDKron --opt NewtonHybrid2PSGDLRA --opt mars-NewtonHybrid2PSGDLRA --opt MSAMLaProp --opt mars-adaptive-NewtonHybrid2PSGDKron  --opt mars-ortho-NewtonHybrid2PSGDKron --opt MuonLaProp --opt mars-unscaled-NewtonHybrid2PSGDKron --opt mars-NewtonHybrid2PSGDKron --opt cautious-AdamW --opt unscaled_cautious-AdamW --opt mars-AdamW  --dtype float32 --steps 1000000 --trials 1000 --parallelism 256 --seeds 1 --difficulties trivial --difficulties easy --difficulties medium --difficulties hard --difficulties extreme --difficulties nightmare --timeout 2880
+> ```
+
+## Migrating from HeavyBall 1.x
+
+- Read the detailed [2.0.0 migration notes](docs/heavyball-2.0.0-migration.md) for an end-to-end checklist, including guidance for restoring legacy behaviour when needed.
+- Use `scripts/migrate_optimizer_state.py` to rewrite pre-2.0 optimizer checkpoints:
+  ```bash
+  python scripts/migrate_optimizer_state.py path/to/checkpoint.pt heavyball.ForeachAdamW --state-key optimizer
+  ```
+  The utility renames legacy state entries, fans them out per parameter view, and injects the HeavyBall metadata block expected by 2.0.0.
+
+
+## Contributing
+
+We welcome contributions! Please check the [issue tracker][tracker] and follow these steps:
+1. Fork the repo and create a feature branch.
+2. Install dev dependencies: `pip install -e .[dev]`.
+3. Run tests: `pytest`.
+4. Submit a pull request.
+
+## License
+
+BSD 3-Clause — see the [LICENSE](LICENSE) file.
+
+---
+<p align="center">
+  Made by the HeavyBall team.
+</p>
+
+[pypi]: https://pypi.org/project/heavyball/
+[license]: LICENSE
+[tracker]: https://github.com/HomebrewML/HeavyBall/issues

heavyball-2.0.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+heavyball/__init__.py,sha256=cabACszT-lZaLPkN2rANFxLxkpcDPb8q4p7iL8jfTtQ,28274
+heavyball/chainable.py,sha256=l7uzXKMbJxKn-kgMR9In8BxYKuvIO_y8uL6P5b0LZo8,41250
+heavyball/helpers.py,sha256=zk_S84wpGcvO9P6kn4UeaQUIDowHxcbM9qQITEm2g5I,30267
+heavyball/utils.py,sha256=sK85OOhmPHvAbUjZhkgX5tRfE3ECai0Yx4Zvt4p8z2Q,97794
+heavyball-2.0.0.dist-info/licenses/LICENSE,sha256=G9fFZcNIVWjU7o6Pr_4sJBRCNDU5X-zelSxIJ2D48ms,1323
+heavyball-2.0.0.dist-info/METADATA,sha256=9JKPFmnvaMT9oigMckzaqI5DjUN0hJj3ePzc1ok3Ia8,4973
+heavyball-2.0.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+heavyball-2.0.0.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
+heavyball-2.0.0.dist-info/RECORD,,

{heavyball-1.7.1.dist-info → heavyball-2.0.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.0)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

heavyball/optimizations/__init__.py

@@ -1,38 +0,0 @@
-"""
-PSGD optimization module - optimized implementations of PSGD functions
-to improve execution speed while maintaining numerical equivalence.
-"""
-
-# Import optimized functions
-# Import integrator API
-from .integrator import (
-    enable_optimizations,
-    get_optimization_status,
-    restore_original_functions,
-)
-from .optimizations import (
-    # LRA optimizations
-    low_rank_mm_optimized,
-    lra_precond_optimized,
-    precond_grad_cached_optimized,
-    # KRON optimizations
-    psgd_calc_A_and_conjB_optimized,
-    psgd_precond_grad_optimized,
-    psgd_update_precond_optimized,
-    update_lra_precond_optimized,
-)
-
-__all__ = [
-    # Optimized functions
-    "low_rank_mm_optimized",
-    "update_lra_precond_optimized",
-    "lra_precond_optimized",
-    "psgd_calc_A_and_conjB_optimized",
-    "psgd_update_precond_optimized",
-    "psgd_precond_grad_optimized",
-    "precond_grad_cached_optimized",
-    # Integrator API
-    "enable_optimizations",
-    "restore_original_functions",
-    "get_optimization_status",
-]

heavyball/optimizations/integrator.py

@@ -1,169 +0,0 @@
-"""
-Integration module to selectively enable optimized implementations
-of PSGD functions while maintaining API compatibility.
-"""
-
-import os
-import sys
-from typing import Any, Dict
-
-import torch
-
-from . import optimizations
-from .. import utils
-
-# Store original function references
-_original_functions = {}
-_optimized_functions = {}
-
-# Mapping of original functions to their optimized versions
-OPTIMIZATION_MAP = {
-    # LRA functions
-    utils.update_lra_precond_: optimizations.update_lra_precond_optimized,
-    utils.lra_precond: optimizations.lra_precond_optimized,
-    # KRON functions
-    utils.psgd_update_precond: optimizations.psgd_update_precond_optimized,
-    utils.psgd_precond_grad: optimizations.psgd_precond_grad_optimized,
-    utils.precond_grad_cached_: optimizations.precond_grad_cached_optimized,
-}
-
-# Config for enabling/disabling optimizations
-_config = {
-    "enabled": os.environ.get("HEAVYBALL_OPTIMIZE", "1") == "1",
-    "torch_compile_allowed": os.environ.get("HEAVYBALL_USE_COMPILE", "1") == "1",
-    "enable_lra": True,
-    "enable_kron": True,
-    "verbose": os.environ.get("HEAVYBALL_VERBOSE", "0") == "1",
-}
-
-
-def _apply_monkey_patch(original_func, optimized_func):
-    """Monkey patch a function with its optimized version."""
-    if original_func not in _original_functions:
-        _original_functions[original_func] = original_func
-
-    # Store reference to the optimized function
-    _optimized_functions[original_func] = optimized_func
-
-    # Get the module where the original function is defined
-    module = original_func.__module__
-    func_name = original_func.__name__
-
-    # Replace the function in its module
-    if hasattr(sys.modules[module], func_name):
-        setattr(sys.modules[module], func_name, optimized_func)
-
-        if _config["verbose"]:
-            print(f"Replaced {module}.{func_name} with optimized version")
-    else:
-        if _config["verbose"]:
-            print(f"Warning: Could not find {func_name} in module {module}")
-
-
-def enable_optimizations(
-    enable: bool = True, lra: bool = True, kron: bool = True, torch_compile: bool = True, verbose: bool = False
-):
-    """
-    Enable or disable PSGD optimizations.
-
-    Args:
-        enable: Whether to enable optimizations at all
-        lra: Whether to enable LRA-specific optimizations
-        kron: Whether to enable Kron-specific optimizations
-        torch_compile: Whether to allow torch.compile optimizations
-        verbose: Whether to print optimization status messages
-    """
-    _config["enabled"] = enable
-    _config["enable_lra"] = lra
-    _config["enable_kron"] = kron
-    _config["torch_compile_allowed"] = torch_compile
-    _config["verbose"] = verbose
-
-    if verbose:
-        print(f"PSGD Optimizations: {'enabled' if enable else 'disabled'}")
-        print(f"  - LRA optimizations: {'enabled' if lra else 'disabled'}")
-        print(f"  - KRON optimizations: {'enabled' if kron else 'disabled'}")
-        print(f"  - torch.compile: {'allowed' if torch_compile else 'disabled'}")
-
-    if not enable:
-        # Restore original functions
-        restore_original_functions()
-        return
-
-    # Apply optimizations based on config
-    for orig_func, opt_func in OPTIMIZATION_MAP.items():
-        # Skip LRA functions if disabled
-        if not _config["enable_lra"] and orig_func in [utils.update_lra_precond_, utils.lra_precond]:
-            continue
-
-        # Skip KRON functions if disabled
-        if not _config["enable_kron"] and orig_func in [
-            utils.psgd_update_precond,
-            utils.psgd_precond_grad,
-            utils.precond_grad_cached_,
-        ]:
-            continue
-
-        _apply_monkey_patch(orig_func, opt_func)
-
-    # Disable torch.compile if not allowed
-    if not _config["torch_compile_allowed"]:
-        # Monkey patch torch.compile to be a no-op
-        def _noop_compile(fn, **kwargs):
-            return fn
-
-        if not hasattr(torch, "_original_compile"):
-            torch._original_compile = torch.compile
-            torch.compile = _noop_compile
-            if verbose:
-                print("Disabled torch.compile (replaced with no-op)")
-    else:
-        # Restore original torch.compile
-        if hasattr(torch, "_original_compile"):
-            torch.compile = torch._original_compile
-            del torch._original_compile
-            if verbose:
-                print("Restored original torch.compile")
-
-
-def restore_original_functions():
-    """Restore all original function implementations."""
-    for orig_func, func_ref in _original_functions.items():
-        module = orig_func.__module__
-        func_name = orig_func.__name__
-
-        if hasattr(sys.modules[module], func_name):
-            setattr(sys.modules[module], func_name, func_ref)
-
-            if _config["verbose"]:
-                print(f"Restored original implementation of {module}.{func_name}")
-
-    # Also restore torch.compile if it was modified
-    if hasattr(torch, "_original_compile"):
-        torch.compile = torch._original_compile
-        del torch._original_compile
-        if _config["verbose"]:
-            print("Restored original torch.compile")
-
-
-def get_optimization_status() -> Dict[str, Any]:
-    """Get current optimization status."""
-    return {
-        "enabled": _config["enabled"],
-        "lra_enabled": _config["enable_lra"],
-        "kron_enabled": _config["enable_kron"],
-        "torch_compile_allowed": _config["torch_compile_allowed"],
-        "optimized_functions": list(_optimized_functions.keys()),
-        "original_functions": list(_original_functions.keys()),
-    }
-
-
-# Auto-initialize optimizations based on environment
-if os.environ.get("HEAVYBALL_AUTO_OPTIMIZE", "1") == "1":
-    enable_optimizations(
-        enable=_config["enabled"],
-        lra=_config["enable_lra"],
-        kron=_config["enable_kron"],
-        torch_compile=_config["torch_compile_allowed"],
-        verbose=_config["verbose"],
-    )

heavyball/optimizations/optimizations.py

@@ -1,329 +0,0 @@
-import random
-from typing import List, Optional
-
-import torch
-from torch import Tensor
-
-from .. import utils
-from ..utils import decorator, decorator_knowngood, min_dtype, scalar_guard, tiny_bf16
-
-#############################
-# PSGD LRA OPTIMIZATIONS
-#############################
-
-
-@decorator
-def low_rank_mm_optimized(U: Tensor, V: Tensor, x: Tensor) -> Tensor:
-    """Optimized version of low_rank_mm using fused operations and memory reuse"""
-    dtype = min_dtype([U, V, x])
-    # Convert only once and cache the result
-    U_dt, V_dt, x_dt = U.to(dtype), V.to(dtype), x.to(dtype)
-
-    # Use a more efficient implementation that avoids multiple conversions
-    # torch.bmm can be more efficient than einsum for this specific pattern
-    if U.dim() == 2:  # This is the common case (batch, rank)
-        # Shape of result: (batch, )
-        tmp = torch.mul(U_dt, x_dt.unsqueeze(-1)).sum(dim=0)  # (rank, )
-        result = torch.mv(V_dt, tmp)  # (batch, )
-        return result.to(x.dtype) + x
-    else:
-        # Fallback to original implementation for other dimensionalities
-        return x + torch.einsum("br,gr,g->b", U_dt, V_dt, x_dt).to(x.dtype)
-
-
-@torch.compile(mode="reduce-overhead")
-def update_lra_precond_core(
-    U: Tensor, V: Tensor, d: Tensor, vector: Tensor, hessian_vector: Tensor, eps: float, step: float, delayed: bool
-):
-    """Core computational part of update_lra_precond optimized with torch.compile"""
-    # Here we apply torch.compile to the computational bottleneck
-    # All inputs are already properly typed and processed
-
-    Qh = low_rank_mm_optimized(U, V, d * hessian_vector)
-    Ph = d * low_rank_mm_optimized(V, U, Qh)
-    rank = U.size(1)
-
-    # Cache VtU computation which is used multiple times
-    VtU = torch.einsum("br,bn->rn", V, U)  # (rank, rank)
-    I = torch.eye(rank, dtype=VtU.dtype, device=VtU.device)
-    IpVtU = I + VtU
-    invQtv = vector / d
-
-    # LU factorization to reuse computation
-    LU, pivots = torch.linalg.lu_factor(IpVtU)
-
-    # Compute vectors inline to reduce memory allocation
-    invQtv = invQtv - V @ torch.linalg.lu_solve(LU, pivots, (U.T @ invQtv).view(-1, 1), adjoint=True).flatten()
-    invPv = invQtv - U @ torch.linalg.lu_solve(LU, pivots, (V.T @ invQtv).view(-1, 1)).flatten()
-    invPv = invPv / d
-
-    # Compute nabla D
-    nablaD = Ph * hessian_vector - vector * invPv
-
-    # Compute divisor more efficiently using fused operations
-    Ph_squared = Ph.square()
-    vector_squared = vector.square()
-    hv_squared = hessian_vector.square()
-    invPv_squared = invPv.square()
-
-    divisor = (Ph_squared + vector_squared) * (hv_squared + invPv_squared)
-    divisor = divisor.add(eps).sqrt().max()
-    d_step = step / divisor
-
-    # Compute for gradient update
-    a, b = Qh, invQtv
-
-    # Update either U or V, not both at the same time
-    precond_u = random.random() < 0.5
-    precond = V if precond_u else U
-
-    # Cache computations that get reused
-    atV = torch.einsum("b,br->r", a, precond)
-    btV = torch.einsum("b,br->r", b, precond)
-    atVVt = torch.einsum("r,br->b", atV, precond)
-    btVVt = torch.einsum("r,br->b", btV, precond)
-
-    # Compute step size
-    precond_step = step / (a.norm() * atVVt.norm() + b.norm() * btVVt.norm() + eps)
-
-    # Update precond matrix
-    if precond_u:
-        a_new = torch.einsum("b,r,rg->bg", a, atV, IpVtU)
-        b_new = torch.einsum("b,r,rg->bg", b, btV, IpVtU)
-    else:
-        # Optimize with in-place operations where possible
-        a_new = a + torch.einsum("br,r->b", V, atV)
-        b_new = b + torch.einsum("br,r->b", V, btV)
-        a_new = torch.einsum("b,r->br", a_new, atV)
-        b_new = torch.einsum("b,r->br", b_new, btV)
-
-    # Return updated values
-    return d, nablaD, d_step, U if precond_u else V, b_new - a_new, precond_step, precond_u
-
-
-def update_lra_precond_optimized(
-    U: List[Tensor],
-    V: List[Tensor],
-    d: List[Tensor],
-    vector: Tensor,
-    hessian_vector: Tensor,
-    eps: float,
-    step: float,
-    delayed: bool,
-):
-    """
-    Optimized version of update_lra_precond_ with:
-    1. Reduced memory allocations
-    2. Fused operations
-    3. Torch.compile for core computations
-    4. Better caching of intermediate results
-    """
-    U_orig, V_orig, d_orig = U, V, d
-
-    # Flatten once
-    U_flat, V_flat, d_flat = utils.flatten(U, 1), utils.flatten(V, 1), utils.flatten(d)
-
-    # Convert dtype once
-    dtype = min_dtype([U_flat, V_flat, vector, hessian_vector])
-    U_dt = U_flat.to(dtype)
-    V_dt = V_flat.to(dtype)
-    vector_dt = vector.to(dtype)
-    hv_dt = hessian_vector.to(dtype)
-
-    # Convert scalar once
-    eps_tensor = scalar_guard(eps, vector)
-
-    try:
-        # Run optimized core computation with torch.compile
-        d_flat, nablaD, d_step, precond, update, precond_step, precond_u = update_lra_precond_core(
-            U_dt, V_dt, d_flat, vector_dt, hv_dt, eps, step, delayed
-        )
-
-        # Apply updates efficiently
-        utils.apply_flat_add(d_orig, d_flat * nablaD, -d_step)
-        utils.apply_flat_add(U_orig if precond_u else V_orig, update, precond_step)
-
-        # For immediate updates
-        if not delayed:
-            utils.stochastic_add_([d], [d_flat * nablaD], -d_step)
-            utils.stochastic_add_([U if precond_u else V], [update], precond_step)
-
-        return U_flat.to(U_orig[0].dtype), V_flat.to(V_orig[0].dtype), d_flat.to(d_orig[0].dtype)
-
-    except RuntimeError:
-        # Fallback to original implementation on failure
-        return utils.update_lra_precond_(U, V, d, vector, hessian_vector, eps, step, delayed)
-
-
-@decorator
-def lra_precond_optimized(U, V, d, g):
-    """
-    Optimized version of lra_precond using memory caching and fused operations
-    """
-    # Get the common dtype only once
-    dtype = min_dtype([U, V, d, g])
-
-    # Convert to this dtype once
-    U_dt, V_dt, d_dt, g_dt = U.to(dtype), V.to(dtype), d.to(dtype), g.to(dtype)
-
-    # First part: g_mid = d * g
-    g_mid = d_dt * g_dt
-
-    # Second part: Qh = low_rank_mm(U, V, g_mid)
-    # Use optimized low_rank_mm
-    Qh = low_rank_mm_optimized(U_dt, V_dt, g_mid)
-
-    # Third part: result = d * low_rank_mm(V, U, Qh)
-    result = d_dt * low_rank_mm_optimized(V_dt, U_dt, Qh)
-
-    # Return result in original dtype
-    return result.to(g.dtype)
-
-
-#############################
-# PSGD KRON OPTIMIZATIONS
-#############################
-
-
-@decorator
-def psgd_calc_A_and_conjB_optimized(exprA, G, Q, conjB):
-    """Optimized version of psgd_calc_A_and_conjB using torch.compile and memory reuse"""
-    order = G.dim()
-    if order > 1:
-        conjB = conjB.view_as(G).permute(*range(1, order), 0)
-
-    # Convert dtype once
-    G_dtype = utils.promote(G.dtype)
-    conjB = conjB.to(G_dtype)
-
-    # Compute A using einsum (could be cached if called multiple times with same Q, G)
-    A = utils.casted_einsum(exprA, *Q, G)
-
-    # Process each Q matrix with potential optimizations
-    for i, q in enumerate(Q):
-        q = utils.promote(q)
-        if q.dim() <= 1:
-            # Scalar case - use in-place division
-            conjB.div_(q)
-        else:
-            # Matrix case - use optimized triangular solve
-            # Reshape once and contiguous to optimize memory access
-            conjB_reshaped = conjB.reshape(-1, q.size(0)).contiguous()
-            solved = torch.linalg.solve_triangular(q, conjB_reshaped, upper=True, left=False)
-            conjB = solved.reshape_as(conjB)
-
-        # Only transpose if needed for next iteration
-        if i < order - 1:
-            conjB = conjB.transpose(i, -1)
-
-    return A, conjB
-
-
-@torch.compile(mode="reduce-overhead")
-def psgd_update_precond_core(Q, term1, term2, precond_lr, norm, q):
-    """Core computation of psgd_update_precond optimized with torch.compile"""
-    term1 *= precond_lr
-    if q.dim() < 2:
-        term1 *= q / norm.clamp_(min=tiny_bf16)
-    else:
-        torch.triu(term1, out=term1)
-        term1 /= torch.where(norm > 0, utils.psgd_lb(term2, norm), norm).clamp_(tiny_bf16)
-        term1 = torch.mm(term1, q)
-    return term1
-
-
-def psgd_update_precond_optimized(Q, exprs, G, precond_lr, oq, store_triu_as_line, V):
-    """Optimized version of psgd_update_precond with reduced allocations and torch.compile"""
-    exprA, exprGs, _ = exprs
-
-    # Use optimized A and conjB calculation
-    A, conjB = psgd_calc_A_and_conjB_optimized(exprA, G, Q, V)
-
-    # Process each Q matrix with optimizations
-    for q, exprG, o in zip(Q, exprGs, oq):
-        # Use optimized einsum implementations
-        term1 = utils.promote(torch.einsum(exprG, A, A))
-        term2 = utils.promote(torch.einsum(exprG, conjB, conjB))
-
-        # Compute the update using compiled core function
-        term1, term2 = term1 - term2, term1 + term2
-        norm = term2.norm(float("inf"))
-
-        try:
-            # Try to use the optimized core calculation
-            term1 = psgd_update_precond_core(Q, term1, term2, precond_lr, norm, q.to(term1.dtype))
-        except (RuntimeError, TypeError):
-            # Fallback to original implementation
-            term1 *= precond_lr
-            if q.dim() < 2:
-                term1 *= q.to(term1.dtype) / norm.clamp_(min=tiny_bf16)
-            else:
-                torch.triu(term1, out=term1)
-                term1 /= torch.where(norm > 0, utils.psgd_lb(term2, norm), norm).clamp_(tiny_bf16)
-                term1 = torch.mm(term1, q.to(term1.dtype))
-
-        # Convert to line format if needed
-        if store_triu_as_line:
-            term1 = utils.triu_to_line([term1])[0][1]
-            # Apply update directly
-            if o.dim() > 0:
-                o.add_(term1)
-            else:
-                o = term1
-        else:
-            # Apply update directly
-            o.add_(term1)
-
-
-@decorator_knowngood
-def psgd_precond_grad_optimized(
-    expr: str, ea: Tensor, *preconds: Tensor, caution: bool = False, grad: Optional[Tensor] = None
-):
-    """Optimized version of psgd_precond_grad with better memory management"""
-    if caution:
-        ea = utils._compilable_cautioning(grad, ea)
-
-    # Determine minimum dtype once
-    md = min_dtype(list(preconds) + [ea])
-
-    # Convert all tensors to the same dtype once
-    args = [q.to(md) for q in preconds]
-    ea_md = ea.to(md)
-
-    # Optimize the einsum operation by avoiding duplicate conversions
-    # and potentially making args contiguous if beneficial
-    args_contiguous = [arg.contiguous() if not arg.is_contiguous() else arg for arg in args]
-    args_double = args_contiguous + args_contiguous
-
-    # Call einsum once with the combined args list
-    new = torch.einsum(expr, *(args_double + [ea_md]))
-
-    # Convert result back to original dtype
-    return new.to(ea.dtype)
-
-
-@decorator_knowngood
-def precond_grad_cached_optimized(
-    expr: str, ea: Tensor, *cached_q: Tensor, caution: bool = False, grad: Optional[Tensor] = None, cast: bool = True
-):
-    """Optimized version of precond_grad_cached_ with better memory management"""
-    if caution:
-        ea = utils._compilable_cautioning(grad, ea)
-
-    # Determine minimum dtype once
-    md = min_dtype(list(cached_q) + [ea])
-
-    # Convert all tensors to the same dtype once and make contiguous if needed
-    args = [q.to(md).contiguous() if not q.is_contiguous() else q.to(md) for q in cached_q]
-    ea_md = ea.to(md).contiguous() if not ea.is_contiguous() else ea.to(md)
-
-    # Add ea_md to args
-    args.append(ea_md)
-
-    # Call einsum once with the optimized args
-    new = torch.einsum(expr, *args)
-
-    # Convert result back if needed
-    if cast:
-        return new.to(ea.dtype)
-    return new