[BUG] Complex vjps are wrong for cos, sqrt, and rsqrt (missing conjugate)
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 28.5k
- Forks
- 2.3k
- Avg merge
- 3d 8h
- Merged PRs (30d)
- 62
Description
Describe the bug
mx.cos, mx.sqrt, and mx.rsqrt return the wrong gradient for complex inputs. cos's vjp delegates to its jvp, and sqrt/rsqrt form the derivative without conjugating it, so all three return f'(z) * cotangent instead of conj(f'(z)) * cotangent (the convention used by mx.exp/mx.log).
To Reproduce
import mlx.core as mx
z = mx.array(0.7 + 0.5j, mx.complex64)
g = mx.array(1.0 + 1.0j, mx.complex64)
print(mx.vjp(mx.cos, [z], [g])[1][0])
print(g * mx.conj(-mx.sin(z))) # expected
Expected behavior
vjp == cotangent * conj(f'(z)), consistent with the other unary ops and finite differences.
Desktop
- OS: macOS
- Version: main (0.32.0.dev)
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start from the mx.cos, mx.sqrt, and mx.rsqrt autodiff entry points and compare their VJP/JVP behavior with mx.exp and mx.log. Reproduce the complex-input case from the issue, then verify that each VJP matches cotangent multiplied by the conjugate derivative and agrees with finite differences.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 68/100