ml-explore / ml-explore/mlx

[BUG] Complex vjps are wrong for cos, sqrt, and rsqrt (missing conjugate)

Open
#3,780 0 comments 0 reactions 0 assignees View on GitHub

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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.