JuliaDiff / JuliaDiff/DifferentiationInterface.jl

Bug(Mooncake): jacobian for real→complex functions returns Float64 instead of ComplexF64

Open
#1,010 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bot
Dominant language
Julia
Stars
313
Forks
35
PR merge metrics
No merged PRs in 30d

Description

Description

When using AutoMooncake() with functions that have real input and complex output, jacobian returns a Float64 matrix containing only partial information instead of the expected ComplexF64 matrix. The imaginary components of the derivatives are lost.

Note: This is likely a Mooncake limitation that should be upstreamed, as the native Mooncake API exhibits the same behavior.

MWE

using DifferentiationInterface
using Mooncake: Mooncake
using ADTypes: AutoMooncake

backend = AutoMooncake(; config=nothing)

# Real input → Complex output
f(x) = [complex(x[1], x[2]), complex(x[2], x[1])]
x = [1.0, 2.0]

J = jacobian(f, backend, x)

Output:

2×2 Matrix{Float64}:
 1.0  0.0
 0.0  1.0

Expected:

2×2 Matrix{ComplexF64}:
 1.0+0.0im  0.0+1.0im
 0.0+1.0im  1.0+0.0im

The jacobian should have imaginary components (the derivative of x[2] appearing in the imaginary part of y[1] is im, not 0).

Native Mooncake API Comparison

The native Mooncake API returns the same incomplete result:

using Mooncake

f(x) = [complex(x[1], x[2]), complex(x[2], x[1])]
x = [1.0, 2.0]

cache = Mooncake.prepare_pullback_cache(f, x)
dy = ComplexF64[1.0, 0.0]

y, (_, dx) = Mooncake.value_and_pullback!!(cache, dy, f, x)
# dx = [1.0, 0.0]  (Float64, not ComplexF64)
# Expected: [1.0, im] or similar complex representation

Mooncake's tangent type appears to match the input type (Float64) rather than accounting for the complex output, losing the imaginary gradient components.

Enzyme Comparison (Correct Output)

Enzyme handles this case correctly:

using DifferentiationInterface
using Enzyme: Enzyme
using ADTypes: AutoEnzyme

f(x) = [complex(x[1], x[2]), complex(x[2], x[1])]
x = [1.0, 2.0]

J = jacobian(f, AutoEnzyme(), x)
# 2×2 Matrix{ComplexF64}:
#  1.0+0.0im  0.0+1.0im
#  0.0+1.0im  1.0+0.0im

Interesting Note

Mooncake's derivative for scalar real→complex works correctly:

g(t) = complex(t, 2*t)
d = derivative(g, AutoMooncake(; config=nothing), 1.0)
# Returns: 1.0 + 2.0im ✓

Only jacobian (which uses pullbacks internally) exhibits this issue.

Recommendation

This should likely be reported upstream to Mooncake, as DI correctly calls the Mooncake API which returns limited results. Possible solutions:

  1. Mooncake could support complex tangents for real inputs when output is complex
  2. DI could document this as a known Mooncake limitation
  3. DI could detect this case and warn users

Environment

  • Julia 1.12.5
  • DifferentiationInterface v0.7.18
  • Mooncake v0.5.28

🤖 I am a robot. This is an experiment in agentic bug-catching under the supervision of @adrhill and @gdalle (#1008). Contents may be hallucinated.

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

Reproduce the real-input/complex-output MWE with jacobian and AutoMooncake, then compare it with derivative and Mooncake's prepare_pullback_cache/value_and_pullback!! entry points. Trace whether DifferentiationInterface transforms or merely returns Mooncake's tangent, and determine whether a DI change is possible or the limitation belongs upstream. Done means a confirmed fix, warning, or documented limitation with regression coverage if this repository owns the behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
backend-api-design
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.