Adam and SGD partially update parameters before rejecting a short gradient list

Bug description

While testing a two-parameter optimization workflow on current main, I found that Adam.step() and SGD.step() do not validate the number of gradient arrays before updating parameters.

With two parameter arrays and one gradient, the first parameter (and applicable optimizer state) is updated before IndexError is raised for the second parameter. The timestep stays at zero. Retrying after correcting the gradient list therefore starts from partially modified state. With three gradients, the extra gradient is silently ignored.

This was found in a small local evaluation, not a production deployment. Investigation and patch preparation were AI-assisted.

Minimal CPU reproduction

import warp as wp
from warp.optim import Adam, SGD

for cls in (Adam, SGD):
    params = [
        wp.full(2, 1.0, device="cpu"),
        wp.full(2, 2.0, device="cpu"),
    ]
    opt = cls(params)
    before = [p.numpy().tolist() for p in params]
    try:
        opt.step([wp.full(2, 0.5, device="cpu")])
    except IndexError as exc:
        print(cls.__name__, type(exc).__name__)
    print("before:", before)
    print("after:", [p.numpy().tolist() for p in params], "t:", opt.t)

Observed:

• Adam first parameter changes from 1.0 to approximately 0.999; second remains 2.0; t == 0.
• SGD first parameter changes from 1.0 to approximately 0.9995; second remains 2.0; t == 0.

Expected behavior and proposed fix

Reject a gradient-count mismatch with ValueError before launching any parameter update. I have a focused local patch adding the count check to both optimizers and regression tests for zero, too few, and too many gradients. The tests check that parameters, moment/momentum buffers, and timestep remain unchanged on rejection.

This proposal addresses gradient-count mismatches only; it does not promise atomicity for every possible invalid gradient or kernel failure.

Validation

• Upstream commit: 85e3eea2f7bfb48649d35e848144990efd0e0a3f.
• Warp 1.19.0.dev0, Linux x86_64, Python 3.12.14, NumPy 2.5.3.
• Built with uv run build_lib.py --no-cuda --no-use-libmathdx.
• Both new test methods fail before the fix (six failing/error subcases).
• After the fix: uv run -m unittest warp.tests.test_adam warp.tests.test_sgd runs 19 tests: 17 pass, 2 CUDA-only tests skipped.
• No NVIDIA GPU was available. CUDA validation is outstanding; no GPU performance or correctness claim is made.

I searched existing issues and PRs for optimizer gradient-count/length validation and did not find the same report. Is this focused validation change welcome?