LAprop as an optimizer along with frobenius weight normalization + associated tests

[mkhona-nvidia]

Convert the LAProp update function into a scalar optimizer

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[skyw]

Shall we seek more programmable way to make optimizer out of scalar update function? Instead of adding them one by one (with a lot of duplicated code) when we need it?

[greptile-apps]

Greptile Summary

This PR introduces LaProp as a new scalar optimizer class, wrapping the pre-existing calculate_laprop_update update function, and adds an optional Frobenius weight normalization feature alongside a full test suite.

• New optimizer class (laprop.py): follows the same WeightDecayMixin + torch.optim.Optimizer pattern as Lion, with per-step lazy state initialization, bias-corrected first/second moment EMAs, and an optional frob_normalize flag that rescales each parameter's Frobenius norm back to its pre-update value.
• Tests (test_laprop.py): cover smoke execution, state initialization, delegate-to-update-function correctness, norm preservation, and all constructor validation errors; the PR also registers laprop in test_registry.py.

Confidence Score: 3/5

The optimizer has several unresolved correctness issues raised in the current PR thread that need to be addressed before merging.

The closure is invoked under @torch.no_grad() without torch.enable_grad(), which silently breaks any training loop that passes a closure. The frob_normalize feature silently renders zero-initialized parameters permanently untrainable, and its interaction with various weight_decay_method settings produces undocumented, inconsistent regularization behavior. These issues are actively discussed in the thread but remain unresolved in the current head commit.

emerging_optimizers/scalar_optimizers/laprop.py — the closure call site, the frob_normalize rescaling path, and the weight-decay/frob-normalize interaction all need attention before this is ready to merge.

Important Files Changed

Filename	Overview
emerging_optimizers/scalar_optimizers/laprop.py	New LaProp optimizer with Frobenius weight normalization; has a known closure/enable_grad bug and frob_normalize edge cases being discussed in the PR thread.
tests/test_laprop.py	Test suite for LaProp; covers smoke, state init, update correctness, norm preservation, and validation errors — gradient coverage gaps and formula fragility are being discussed in the PR thread.
emerging_optimizers/scalar_optimizers/init.py	Adds LaProp import in alphabetical order before Lion; no issues.
tests/test_registry.py	Adds laprop registry entry to the parameterized get_optimizer test; straightforward and correct.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[LaProp.step called] --> B{closure?}
    B -- yes --> C[loss = closure
No enable_grad!]
    B -- no --> D[loss = None]
    C --> E[for each param group]
    D --> E
    E --> F[_init_group: lazy init
exp_avg, exp_avg_sq, step=0]
    F --> G[for each param p
with gradient]
    G --> H[step += 1]
    H --> I{frob_normalize?}
    I -- yes --> J[pre_norm = p.data.norm
before weight decay]
    I -- no --> K[pre_norm = None]
    J --> L[_apply_weight_decay_inplace
modifies p.data or grad]
    K --> L
    L --> M[calculate_laprop_update
updates exp_avg_sq, exp_avg in-place
returns bias-corrected EMA]
    M --> N[p.data += -lr * update]
    N --> O{frob_normalize?}
    O -- yes --> P[p.data *= pre_norm /
max p.data.norm, eps]
    O -- no --> Q[next param]
    P --> Q
    Q --> G

Loading

_{Reviews (7): Last reviewed commit: "removed enable grad from closure path" | Re-trigger Greptile}

[greptile-apps]

The normalized_grad formula in the expected-value calculation is only algebraically equivalent to the actual implementation for step 1 with correct_bias=True. At that step, bias_correction2 equals (1 - beta2), which exactly cancels the lerp coefficient, reducing the denominator to |grad| + eps. Passing correct_bias=False (or any step > 1) would produce a different denominator in calculate_laprop_update, but the hardcoded formula here would not change — so the assertion would silently fail for incorrect reasons. Adding an explicit correct_bias=True makes the assumption visible and protects against future regressions.

Suggested change

	optimizer = LaProp([param], lr=0.25, betas=(beta1, beta2), weight_decay=0.0)
	optimizer = LaProp([param], lr=0.25, betas=(beta1, beta2), weight_decay=0.0, correct_bias=True)

[greptile-apps]

Test only exercises positive gradients

torch.randint(1, 5, ...) produces values in [1, 4], so grad.abs() == grad for every element. The normalization formula grad / (grad.abs() + eps) is therefore never tested with negative values, where the sign of the normalized gradient is the critical part of LaProp's behaviour. A torch.randn-generated gradient (or explicitly mixed-sign integers) would provide better coverage.

[skyw]

Nit: Do we need two eps? They are for different things, but I feel having two of them only increase probability to mess up one of them.

[skyw]

I know greptile like to signal this, but all of our optimizer are not supposed to have gradient in it. can just assert closure is None and ignore greptile.

[skyw]

Use self.normalize as condition.

[skyw]

Nit: Consider a more informative name to indicate what normalization the flag enables.

[skyw]

Q: Should they exact match?

[skyw]

Test name suggest "match", both atol and rtol needs to be 0.

[greptile-apps]

When normalize=True and a parameter has been zero-initialized, pre_norm is 0.0. After the LaProp gradient update produces non-zero values, the line below multiplies the whole tensor by 0 / clamp_min(positive, eps) = 0, silently zeroing the parameter back out on every step. The parameter can never leave zero, so the normalization feature renders zero-initialized parameters completely untrainable without any warning.

Suggested change

	if pre_norm is not None:
	p.data.mul_(pre_norm / p.data.norm().clamp_min(self.normalize_eps))
	if pre_norm is not None and pre_norm >= self.normalize_eps:
	p.data.mul_(pre_norm / p.data.norm().clamp_min(self.normalize_eps))

[skyw]

We should do a full repo clean up for the closure. We don't support it in any of our optimizer, i.e. none was tested, should just error out if closure is not none.

[mkhona-nvidia]

/ok to test 9f2c49e

[mkhona-nvidia]

/ok to test 9ab27ce

[mkhona-nvidia]

/ok to test d6c0cee