Repository navigation
[MLIR] fix: complex.create reverse mode flips imaginary-operand gradient sign - #2899
Merged
Merged
Conversation
f(t) = im(create(0, t)) = t is a smooth real identity, so df/dt must be +1. The existing complex.mlir covers re/im standalone and create in the forward position, but not the consumer path where create feeds im. Adds that case; asserts no arith.negf survives (im's adjoint negation cancels the Wirtinger conj wrap when the gradient is correct).
The prior CHECK-NOT: arith.negf was a poor discriminator: im's adjoint legitimately emits a negf even when correct. Assert the actual return is the positive seed (df/dt=+1) so a failure pins the sign flip and dumps the return.
The MLIR reverse wrap conjugates the incoming cotangent of every complex-result op (enzyme-tblgen.cpp, createConjOp on diffe). For holomorphic ops this cancels against the per-contribution conj; for complex.create (R^2->C) it does not: the rule extracted ComplexIm(DiffeRet) = Im(conj(z_bar)) = -Im(z_bar), negating the gradient flowing to the imaginary operand. So f(t) = im(create(0,t)) = t differentiated to -1 instead of +1. Negate the imaginary-operand contribution to compensate, mirroring the existing NegF compensation in the ImOp rule. Only complex.create in primal position is affected (no existing lit test exercises it; complex.mlir uses create only in the adjoint). Covered by complex_create_im.mlir.
kimjune01
changed the title
CI confirmed the fix flips the gradient to +seed, but canonicalize does not fold the corrected chain to 'return %arg1' (it leaves the negf/conj chain and returns %8 = +seed). Assert the fix's signature instead: a negf of the complex.im result. Without the fix, complex.im is the returned cotangent (-seed) with no following negf, so this still discriminates the bug.
kimjune01
marked this pull request as ready for review
June 25, 2026 00:05
mofeing
approved these changes
Jun 26, 2026
mofeing
left a comment
mofeing
left a comment
Collaborator
There was a problem hiding this comment.
thanks! this is correct. we do the same thing in Enzyme-JAX for the stablehlo::ComplexOp op.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #2900.
complex.createin reverse mode negates the gradient on its imaginary operand. The reverse generator conjugates the cotangent it fetches for each complex-result op (createConjOpon the diffe), and the other complex rules account for that;CreateOpdoes not, so its rule readsComplexIm(DiffeRet) = Im(conj(z_bar)) = -Im(z_bar). Its contributions are realf32(wherecreateConjOpis identity), so nothing downstream compensates. Full reasoning and the emitted-IR receipt are in #2900.Change
CreateOprule, matching theNegFtheImOprule already uses, soCreateOpis consistent with the cotangent convention the other complex rules follow.complex.reis unaffected. Onlycomplex.createin primal position changes;complex.mlirusescreateonly in the adjoint, so nothing there moves.test/MLIR/ReverseMode/complex_create_im.mlir: differentiatesf(t) = im(create(0, t)) = tand checks the returned cotangent is+seed. The corrected adjoint chain is an identity, so it canonicalizes toreturn %arg1; the buggy-seedchain cannot.Verification
Developed against your CI (no local MLIR build here). The test-only commit's MLIR run went red on the sign, which is the bug. The post-fix run currently needs maintainer approval to start (
action_required), so a green-after-fix is pending your approval rather than confirmed by me. Happy to adjust a CHECK line against the exact canonicalized IR once it runs.