Skip to content

Add Enzyme rules for factorizations - #464

Draft
kshyatt wants to merge 4 commits into
mainfrom
ksh/enz_fact
Draft

kshyatt wants to merge 4 commits into
mainfrom
ksh/enz_fact

Conversation

@kshyatt

@kshyatt kshyatt commented Jun 26, 2026

Copy link
Copy Markdown
Member

For some reason svd_compact and svd_trunc_no_error don't play nicely with Enzyme here using the MAK rules, I think because of the DiagonalTensorMap output. I can try to investigate further if preferred. I also added some additional logic in pullbacks to match what the MAK pullbacks actually kick back.

@kshyatt
kshyatt requested a review from lkdvos June 26, 2026 00:48
@kshyatt

kshyatt commented Jun 26, 2026

Copy link
Copy Markdown
Member Author

Oh crap lol I was working on top of a dev-ed MAK

@codecov

codecov Bot commented Jun 26, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 82.00000% with 9 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/factorizations/pullbacks.jl 71.42% 6 Missing ⚠️
ext/TensorKitEnzymeExt/factorizations.jl 96.29% 1 Missing ⚠️
ext/TensorKitEnzymeExt/utility.jl 0.00% 1 Missing ⚠️
src/factorizations/diagonal.jl 0.00% 1 Missing ⚠️
Files with missing lines Coverage Δ
ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl 100.00% <ø> (ø)
ext/TensorKitEnzymeExt/factorizations.jl 96.29% <96.29%> (ø)
ext/TensorKitEnzymeExt/utility.jl 20.45% <0.00%> (-0.24%) ⬇️
src/factorizations/diagonal.jl 67.64% <0.00%> (-2.05%) ⬇️
src/factorizations/pullbacks.jl 70.76% <71.42%> (+0.31%) ⬆️

... and 3 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@kshyatt
kshyatt marked this pull request as draft June 26, 2026 14:10
@kshyatt

kshyatt commented Jun 26, 2026

Copy link
Copy Markdown
Member Author

Weirdly the eig failures seem to go away on MatrixAlgebraKit master, let me try with some [sources] and see if it's reproducible outside of my laptop

@github-actions

github-actions Bot commented Jun 29, 2026 •

Copy link
Copy Markdown
Contributor

Your PR no longer requires formatting changes. Thank you for your contribution!

@kshyatt
kshyatt force-pushed the ksh/enz_fact branch 4 times, most recently from e0a9f0c to 9742ab7 Compare July 6, 2026 18:14
@kshyatt
kshyatt force-pushed the ksh/enz_fact branch 3 times, most recently from 049c1a0 to a60a6ca Compare July 15, 2026 14:21
@kshyatt

kshyatt commented Jul 16, 2026

Copy link
Copy Markdown
Member Author

This map business with block is not the nicest but I thought it was better than letting block(nothing, c) go through silently in other cases

return Δt
end
@eval function MAK.$pullback!(
Δt::AbstractTensorMap, ::Nothing, F, ΔF, inds; kwargs...

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Out of curiosity, who is generating the nothings here?

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

nvm, I see, do we also have this in the other parts of the code?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

What do you mean by other parts? It happens for pbs where the value of t doesn't contribute, it's a pattern inherited from MAK.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

thanks, that was exactly the answer I was after :) I guess there we don't have explicit dispatch for them, and use iszerotangent in the function bodies. Is that worth it to do here as well? Would it be as simple as nothing_or_block(x, c) = isnothing(x) ? x : block(x, c)?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

yeah we could do that, i'm always biased towards making another method but being less verbose is good

@kshyatt
kshyatt marked this pull request as draft September 30, 2026 18:26
@kshyatt

kshyatt commented Sep 30, 2026

Copy link
Copy Markdown
Member Author

Gonna make this a draft since 1.10 is going to fail anyway and the forward mode stuff needs some love

This branch has not been deployed

No deployments
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