Skip to content

Fix complex VJPs for log and exp#3605

Merged
zcbenz merged 1 commit into
ml-explore:mainfrom
CameronChurchwell:complex_exp_log_vjp_fix
Jun 4, 2026
Merged

Fix complex VJPs for log and exp#3605
zcbenz merged 1 commit into
ml-explore:mainfrom
CameronChurchwell:complex_exp_log_vjp_fix

Conversation

@CameronChurchwell
Copy link
Copy Markdown
Contributor

Fix complex VJPs for log and exp

This is the same as #3433, just with the exp and log functions now. Specifically, the VJP implementations for these primitives were lacking a conjugation which is necessary for correctness under the Wirtinger convention used by MLX. The change to Log::vjp looks more extensive since the previous implementation re-uses Log::jvp, but I have effectively copied over and modified the body of Log::jvp, so the actual change is just the addition of the conjugation

I have added tests to verify correctness, and I ran pre-commit already. No documentation changes should be necessary, as this is just a correctness fix.

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

Copy link
Copy Markdown
Collaborator

@zcbenz zcbenz left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me, thanks!

@zcbenz zcbenz merged commit b155224 into ml-explore:main Jun 4, 2026
16 checks passed
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.

2 participants