Skip to content

Enzyme and Mooncake fwd mode support#268

Merged
kshyatt merged 3 commits intomasterfrom
ksh/fwd
Apr 29, 2026
Merged

Enzyme and Mooncake fwd mode support#268
kshyatt merged 3 commits intomasterfrom
ksh/fwd

Conversation

@kshyatt
Copy link
Copy Markdown
Member

@kshyatt kshyatt commented Apr 28, 2026

Added forward rules and tests. Might be nice for doing forward mode AD for TNR in the future.

@codecov
Copy link
Copy Markdown

codecov Bot commented Apr 28, 2026

Codecov Report

❌ Patch coverage is 53.21101% with 51 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
...orOperationsEnzymeExt/TensorOperationsEnzymeExt.jl 0.00% 51 Missing ⚠️
Files with missing lines Coverage Δ
...erationsMooncakeExt/TensorOperationsMooncakeExt.jl 96.63% <100.00%> (+3.19%) ⬆️
...orOperationsEnzymeExt/TensorOperationsEnzymeExt.jl 0.00% <0.00%> (ø)
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@kshyatt kshyatt requested a review from Jutho April 29, 2026 07:28
Comment thread ext/TensorOperationsEnzymeExt/TensorOperationsEnzymeExt.jl Outdated
Comment thread ext/TensorOperationsEnzymeExt/TensorOperationsEnzymeExt.jl Outdated
Comment thread ext/TensorOperationsEnzymeExt/TensorOperationsEnzymeExt.jl Outdated
Comment thread ext/TensorOperationsMooncakeExt/TensorOperationsMooncakeExt.jl Outdated
Comment thread ext/TensorOperationsMooncakeExt/TensorOperationsMooncakeExt.jl Outdated
Comment thread ext/TensorOperationsMooncakeExt/TensorOperationsMooncakeExt.jl Outdated
@Jutho
Copy link
Copy Markdown
Member

Jutho commented Apr 29, 2026

Looks good to me. I think these rules can be relaxed to not be restricted to AbstractArray subtypes; they would then immediately work for TensorMap as well.

This then also makes me trying to recall why we restricted the reverse rules? The only "array"-specific operation in most of these rules seems to be Mooncake.arrify, but since TensorMaps support this, would it not simplify our codebase to also relax the reverse rules?

@Jutho
Copy link
Copy Markdown
Member

Jutho commented Apr 29, 2026

Also, these tests are currently not run. TensorOperations doesn't use ParallelTestRunners and the automatic detection stuff yet, so you have to manually include the files in runtests.jl

@kshyatt
Copy link
Copy Markdown
Member Author

kshyatt commented Apr 29, 2026

Wait, aren't all the tests in the test/mooncake.jl and test/enzyme.jl test files which are run?

@lkdvos
Copy link
Copy Markdown
Member

lkdvos commented Apr 29, 2026

I think we had to redo them for the tensormaps because of the additional fermionic signs...

@kshyatt
Copy link
Copy Markdown
Member Author

kshyatt commented Apr 29, 2026

I think the GitHub actions aren't uploading coverage data properly, I opened another PR about it

@Jutho
Copy link
Copy Markdown
Member

Jutho commented Apr 29, 2026

Wait, aren't all the tests in the test/mooncake.jl and test/enzyme.jl test files which are run?

Oh yes, I thought these were new files, and only saw the zero% PR coverage message, so I inferred that the tests did not run.

@Jutho
Copy link
Copy Markdown
Member

Jutho commented Apr 29, 2026

I think we had to redo them for the tensormaps because of the additional fermionic signs...

Damn you, fermions.

kshyatt and others added 3 commits April 29, 2026 14:06
Co-authored-by: Jutho <Jutho@users.noreply.github.com>
@kshyatt
Copy link
Copy Markdown
Member Author

kshyatt commented Apr 29, 2026

I've rebased on top of the Codecov PR, let's see if this looks a little better. Unfortunately the Enzyme code will still be uncovered...

@kshyatt
Copy link
Copy Markdown
Member Author

kshyatt commented Apr 29, 2026

There we go, looks like the Codecov is working again

@kshyatt
Copy link
Copy Markdown
Member Author

kshyatt commented Apr 29, 2026

@Jutho does this look good to merge ?

@Jutho
Copy link
Copy Markdown
Member

Jutho commented Apr 29, 2026

Somehow it still seems as if none of the Enzyme changes are covered, whereas 100% of the Mooncake changes are.

@kshyatt
Copy link
Copy Markdown
Member Author

kshyatt commented Apr 29, 2026

Yes, unfortunately due to weird lowering behaviour, the Julia coverage cannot "see" into the Enzyme rules. We see the same behaviour over at MAK.

@kshyatt kshyatt merged commit cb2039f into master Apr 29, 2026
12 of 13 checks passed
@kshyatt kshyatt deleted the ksh/fwd branch April 29, 2026 17:32
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.

3 participants