tests/test_optim.py

import torch

from stego_olmoe.optim import Muon, zeropower_via_newtonschulz5


def test_newton_schulz_returns_finite_matrix():
    update = torch.randn(8, 4)
    out = zeropower_via_newtonschulz5(update, steps=2)

    assert out.shape == update.shape
    assert torch.isfinite(out).all()


def test_muon_updates_matrix_and_vector_params():
    matrix = torch.nn.Parameter(torch.randn(8, 4))
    vector = torch.nn.Parameter(torch.randn(4))
    matrix_before = matrix.detach().clone()
    vector_before = vector.detach().clone()
    matrix.grad = torch.randn_like(matrix)
    vector.grad = torch.randn_like(vector)

    optimizer = Muon([matrix, vector], lr=1e-3)
    optimizer.step()

    assert not torch.equal(matrix.detach(), matrix_before)
    assert not torch.equal(vector.detach(), vector_before)
    assert "momentum_buffer" in optimizer.state[matrix]
    assert "exp_avg" in optimizer.state[vector]


def test_native_torch_muon_rejects_vector_params_when_available():
    if not hasattr(torch.optim, "Muon"):
        return
    matrix = torch.nn.Parameter(torch.randn(4, 4))
    vector = torch.nn.Parameter(torch.randn(4))

    try:
        torch.optim.Muon([matrix, vector], lr=1e-3)
    except ValueError as exc:
        assert "2D" in str(exc)
    else:
        raise AssertionError("native torch.optim.Muon unexpectedly accepted a vector parameter")