Skip to content

Run AdaBound on the real and imaginary parts - #555

Open
shaneraphel wants to merge 1 commit into
jettify:masterfrom
shaneraphel:adabound-complex
Open

shaneraphel wants to merge 1 commit into
jettify:masterfrom
shaneraphel:adabound-complex

Conversation

@shaneraphel

@shaneraphel shaneraphel commented Sep 25, 2026 •

Copy link
Copy Markdown

Summary

Fixes #458 for AdaBound. step_size.clamp_ is not implemented for complex tensors, so a complex weight raised NotImplementedError: clamp is not supported for complex types at the learning-rate bound and never updated.

A complex parameter is updated as two real components. The moments, the bound and the write-back run on torch.view_as_real, which is how Adam updates a complex weight: both the real and the imaginary part are kept, and neither is dropped. Real parameters take the same statements as before.

Adahessian's complex failure in the same issue is a different call (torch.randint_like on a complex tensor) and is not changed here.

Test plan

  • pytest tests/test_adabound_complex.py
  • Two real steps with weight decay match the previous implementation exactly, including the value 0.09872793406248093 on the first coordinate.
  • A complex step, and the AMSBound variant, match AdaBound run on view_as_real of the same parameter and gradient, bitwise.

This change was produced with assistance from an AI coding tool. I reproduced the clamp error, kept both components through a real view, and checked the real path against the previous code.

clamp is not implemented for complex tensors, so a complex weight
raised NotImplementedError at the learning-rate bound (issue 458).
The bound is applied to a real view of the parameter, the same way
Adam updates a complex weight: both components are kept and the real
parameter path is unchanged.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Complex numbers

1 participant