Skip to content

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

Closed

Description

@obchain

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)

Activity

  1. zcbenz commented on Sep 26, 2026

    @zcbenz
    Member

    Merging into #3778

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

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions