Add Enzyme rules and move some logic into pullbacks

[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 151 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...orOperationsEnzymeExt/TensorOperationsEnzymeExt.jl	0.00%	83 Missing ⚠️
src/pullbacks/contract.jl	0.00%	26 Missing ⚠️
src/pullbacks/trace.jl	0.00%	18 Missing ⚠️
src/pullbacks/add.jl	0.00%	14 Missing ⚠️
...erationsMooncakeExt/TensorOperationsMooncakeExt.jl	0.00%	9 Missing ⚠️
src/utils.jl	0.00%	1 Missing ⚠️

Files with missing lines	Coverage Δ
ext/TensorOperationsChainRulesCoreExt.jl	0.00% <ø> (ø)
src/TensorOperations.jl	100.00% <ø> (ø)
src/utils.jl	0.00% <0.00%> (ø)
...erationsMooncakeExt/TensorOperationsMooncakeExt.jl	0.00% <0.00%> (ø)
src/pullbacks/add.jl	0.00% <0.00%> (ø)
src/pullbacks/trace.jl	0.00% <0.00%> (ø)
src/pullbacks/contract.jl	0.00% <0.00%> (ø)
...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.

[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!