Add Enzyme rules

[kshyatt]

We can probably use these PB functions for ChainRules as well though I'm less familiar with that library, so I might need to ask someone else to do the plumbing :)

[codecov]

Codecov Report

❌ Patch coverage is 0% with 78 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...orOperationsEnzymeExt/TensorOperationsEnzymeExt.jl	0.00%	78 Missing ⚠️

Files with missing lines	Coverage Δ
...orOperationsEnzymeExt/TensorOperationsEnzymeExt.jl	0.00% <0.00%> (ø)

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[lkdvos]

I'll gladly update the chainrules implementation once we are happy with the names and signatures etc

[kshyatt]

Since I'm already modifying the Mooncake stuff let's also wait for Jutho/StridedViews.jl#19 to land so that we can test the complex stuff with those rules.

[lkdvos]

I saw you removed the import of the chainrules overloads for tensoralloc and tensorfree. I think we still actually need this, for example in the ManualAllocator case we are both capturing the intermediate tensors in the closures and we might be freeing them at the end of the contraction, so this would lead to a use after free (and possibly segfaults).

For Chainrules, we have taken the cautious approach and simply disabled having temporary tensors and freeing them (which is exactly these overloads, i.e. temp=Val(true) is replaced by temp = Val(false)), and I would think we need this here too?

TensorOperations.jl/ext/TensorOperationsChainRulesCoreExt.jl

Lines 26 to 41 in 22ef6dc

	function ChainRulesCore.rrule(
	::typeof(TensorOperations.tensorfree!), allocator = DefaultAllocator()
	)
	tensorfree!_pullback(Δargs...) = (NoTangent(), NoTangent())
	return nothing, tensorfree!_pullback
	end
	function ChainRulesCore.rrule(
	::typeof(TensorOperations.tensoralloc), ttype, structure,
	istemp, allocator = DefaultAllocator()
	)
	output = TensorOperations.tensoralloc(ttype, structure, Val(false), allocator)
	function tensoralloc_pullback(Δargs...)
	return (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent())
	end
	return output, tensoralloc_pullback
	end

I'm guessing it will be better to start using the shared pullbacks for Mooncake and Chainrules in separate PRs, so probably best to leave that for now as is?

Otherwise very cool!

[kshyatt]

Running the Ubuntu test locally to see if I can reproduce, it's quite odd imo

[kshyatt]

Can't repro the ubuntu failure locally... annoying

[kshyatt]

@lkdvos I think is ready to review. 1.10 isn't working but the Enzyme crew are working on figuring it out, I think we can enable those tests once they have this fixed.

[kshyatt]

Suggestions seem to have created new test failures, I'll see if I can figure out why

[kshyatt]

OK, it looks like the chances to cache_C are responsible, and annoyingly, doing something like

cache_C = !iszero(β_dβ.val) ? copy(C_dC.val) : C_dC.val (for type stability)

doesn't help things. So I might undo the changes but leave some comments to revisit.

[kshyatt]

Seems like things are working now :)

[Jutho]

If all tests pass, I think this is good to go.

[Jutho]

OK, it looks like the chances to cache_C are responsible, and annoyingly, doing something like

cache_C = !iszero(β_dβ.val) ? copy(C_dC.val) : C_dC.val (for type stability)

doesn't help things. So I might undo the changes but leave some comments to revisit.

So this approach did work in the end? Was there another cause for type instabilities then?

[kshyatt]

So this approach did work in the end? Was there another cause for type instabilities then?

This did work, it seems to have fixed compiler issues on 1.12. On 1.10, there's a weird memory corruption bug I was unable to diagnose but the Enzyme devs are digging into it.

[Jutho]

I might have created a merge conflict by merging the compat thing (which jsut reorganized the order of the compat entries).