From dac97cc2b7a06f11a936be4cf5412518367048f3 Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <khyatt@flatironinstitute.org>
Date: Fri, 30 Oct 2026 13:17:53 -0400
Subject: [PATCH 1/2] Add some rules for a tensorfree which also takes the
 tensor

---
 ext/TensorOperationsChainRulesCoreExt.jl                    | 6 ++++++
 ext/TensorOperationsEnzymeExt/TensorOperationsEnzymeExt.jl  | 1 +
 .../TensorOperationsMooncakeExt.jl                          | 1 +
 3 files changed, 8 insertions(+)

diff --git a/ext/TensorOperationsChainRulesCoreExt.jl b/ext/TensorOperationsChainRulesCoreExt.jl
index eeee38a7..a46fc9bf 100644
--- a/ext/TensorOperationsChainRulesCoreExt.jl
+++ b/ext/TensorOperationsChainRulesCoreExt.jl
@@ -33,6 +33,12 @@ function ChainRulesCore.rrule(
     tensorfree!_pullback(Δargs...) = (NoTangent(), NoTangent())
     return nothing, tensorfree!_pullback
 end
+function ChainRulesCore.rrule(
+        ::typeof(TensorOperations.tensorfree!), C, allocator = DefaultAllocator()
+    )
+    tensorfree!_pullback(Δargs...) = (NoTangent(), NoTangent(), NoTangent())
+    return nothing, tensorfree!_pullback
+end
 function ChainRulesCore.rrule(
         ::typeof(TensorOperations.tensoralloc), ttype, structure,
         istemp, allocator = DefaultAllocator()
diff --git a/ext/TensorOperationsEnzymeExt/TensorOperationsEnzymeExt.jl b/ext/TensorOperationsEnzymeExt/TensorOperationsEnzymeExt.jl
index 8094b908..97c92ed6 100644
--- a/ext/TensorOperationsEnzymeExt/TensorOperationsEnzymeExt.jl
+++ b/ext/TensorOperationsEnzymeExt/TensorOperationsEnzymeExt.jl
@@ -9,6 +9,7 @@ using Enzyme.EnzymeCore
 using Enzyme.EnzymeCore: EnzymeRules
 
 @inline EnzymeRules.inactive(::typeof(TensorOperations.tensorfree!), ::Any) = true
+@inline EnzymeRules.inactive(::typeof(TensorOperations.tensorfree!), ::Any, ::Any) = true
 @inline EnzymeRules.inactive_type(v::Type{<:AbstractBackend}) = true
 @inline EnzymeRules.inactive_type(v::Type{DefaultAllocator}) = true
 @inline EnzymeRules.inactive_type(v::Type{<:CUDAAllocator}) = true
diff --git a/ext/TensorOperationsMooncakeExt/TensorOperationsMooncakeExt.jl b/ext/TensorOperationsMooncakeExt/TensorOperationsMooncakeExt.jl
index 57acd534..2ea5db53 100644
--- a/ext/TensorOperationsMooncakeExt/TensorOperationsMooncakeExt.jl
+++ b/ext/TensorOperationsMooncakeExt/TensorOperationsMooncakeExt.jl
@@ -27,6 +27,7 @@ Mooncake.tangent_type(::Type{ManualAllocator}) = Mooncake.NoTangent
 @zero_derivative DefaultCtx Tuple{typeof(TensorOperations.promote_add), Any}
 
 Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(TensorOperations.tensorfree!), Any}
+Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(TensorOperations.tensorfree!), Any, Any}
 Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(TensorOperations.tensoralloc), Any, Any, Any, Any}
 
 Mooncake.@is_primitive DefaultCtx Tuple{typeof(tensorcontract!), AbstractArray, AbstractArray, Index2Tuple, Bool, AbstractArray, Index2Tuple, Bool, Index2Tuple, Number, Number, Vararg{Any}}

From 5f9825f7a37a0d4360cd49ef89ae86ff5b40096d Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <khyatt@flatironinstitute.org>
Date: Fri, 30 Oct 2026 13:22:08 -0400
Subject: [PATCH 2/2] dumb typo

---
 ext/TensorOperationsChainRulesCoreExt.jl | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/ext/TensorOperationsChainRulesCoreExt.jl b/ext/TensorOperationsChainRulesCoreExt.jl
index a46fc9bf..02bdae1b 100644
--- a/ext/TensorOperationsChainRulesCoreExt.jl
+++ b/ext/TensorOperationsChainRulesCoreExt.jl
@@ -27,15 +27,11 @@ using LinearAlgebra
 
 # Cannot free intermediate tensors when using AD
 # Thus we change the forward passes: `istemp=false` and `tensorfree!` is a no-op
-function ChainRulesCore.rrule(
-        ::typeof(TensorOperations.tensorfree!), allocator = DefaultAllocator()
-    )
+function ChainRulesCore.rrule(::typeof(TensorOperations.tensorfree!), C)
     tensorfree!_pullback(Δargs...) = (NoTangent(), NoTangent())
     return nothing, tensorfree!_pullback
 end
-function ChainRulesCore.rrule(
-        ::typeof(TensorOperations.tensorfree!), C, allocator = DefaultAllocator()
-    )
+function ChainRulesCore.rrule(::typeof(TensorOperations.tensorfree!), C, allocator)
     tensorfree!_pullback(Δargs...) = (NoTangent(), NoTangent(), NoTangent())
     return nothing, tensorfree!_pullback
 end