From f4bb1bc52c37e4057216e8b26d946fa524208868 Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <katharine.hyatt@ugent.be>
Date: Mon, 10 Aug 2026 16:47:28 +0200
Subject: [PATCH 1/3] Special case tensoradd! to bypass PermutedDimsArray for
 Diagonals

---
 src/implementation/base.jl | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/src/implementation/base.jl b/src/implementation/base.jl
index ecd2aed9..3e0efae8 100644
--- a/src/implementation/base.jl
+++ b/src/implementation/base.jl
@@ -28,6 +28,32 @@ function tensoradd!(
     end
     return C
 end
+function tensoradd!(
+        C::AbstractArray,
+        A::Diagonal, pA::Index2Tuple, conjA::Bool,
+        α::Number, β::Number,
+        ::BaseView, allocator = DefaultAllocator()
+    )
+    argcheck_tensoradd(C, A, pA)
+    dimcheck_tensoradd(C, A, pA)
+
+    # can we assume that C is mutable?
+    # is there more functionality in base that we can use?
+    if conjA
+        if iszero(β)
+            C .= α .* conj.(A)
+        else
+            C .= β .* C .+ α .* conj.(A)
+        end
+    else
+        if iszero(β)
+            C .= α .* A
+        else
+            C .= β .* C .+ α .* A
+        end
+    end
+    return C
+end
 function tensoradd!(
         C::AbstractArray,
         A::AbstractArray, pA::Index2Tuple, conjA::Bool,

From bb1cee6dec487d12dae93de66b4afbf7fc9aac7f Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <katharine.hyatt@ugent.be>
Date: Mon, 10 Aug 2026 16:52:44 +0200
Subject: [PATCH 2/3] Loosen backend req for Diagonal

---
 src/implementation/base.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/implementation/base.jl b/src/implementation/base.jl
index 3e0efae8..7c903e7f 100644
--- a/src/implementation/base.jl
+++ b/src/implementation/base.jl
@@ -32,7 +32,7 @@ function tensoradd!(
         C::AbstractArray,
         A::Diagonal, pA::Index2Tuple, conjA::Bool,
         α::Number, β::Number,
-        ::BaseView, allocator = DefaultAllocator()
+        backend, allocator = DefaultAllocator()
     )
     argcheck_tensoradd(C, A, pA)
     dimcheck_tensoradd(C, A, pA)

From 2de1bb139d1a401db4b604c7a82b3d50f7f9dcd4 Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <katharine.hyatt@ugent.be>
Date: Tue, 11 Aug 2026 10:02:01 +0200
Subject: [PATCH 3/3] Fix

---
 src/implementation/base.jl | 15 +++++++++++----
 test/gpu.jl                | 26 ++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/src/implementation/base.jl b/src/implementation/base.jl
index 7c903e7f..2410436a 100644
--- a/src/implementation/base.jl
+++ b/src/implementation/base.jl
@@ -28,11 +28,13 @@ function tensoradd!(
     end
     return C
 end
+
+# needed to avoid ambiguity error
 function tensoradd!(
         C::AbstractArray,
         A::Diagonal, pA::Index2Tuple, conjA::Bool,
         α::Number, β::Number,
-        backend, allocator = DefaultAllocator()
+        ::BaseView, allocator = DefaultAllocator()
     )
     argcheck_tensoradd(C, A, pA)
     dimcheck_tensoradd(C, A, pA)
@@ -65,8 +67,13 @@ function tensoradd!(
 
     # can we assume that C is mutable?
     # is there more functionality in base that we can use?
-    Atemp = tensoralloc_add(eltype(A), A, pA, conjA, Val(true), allocator)
-    Ã = permutedims!(Atemp, A, linearize(pA))
+    Ã = if !isa(A, Diagonal)
+        Atemp = tensoralloc_add(eltype(A), A, pA, conjA, Val(true), allocator)
+        permutedims!(Atemp, A, linearize(pA))
+    else
+        Atemp = nothing
+        A
+    end
     if conjA
         if iszero(β)
             C .= α .* conj.(Ã)
@@ -80,7 +87,7 @@ function tensoradd!(
             C .= β .* C .+ α .* Ã
         end
     end
-    tensorfree!(Atemp, allocator)
+    !isa(A, Diagonal) && tensorfree!(Atemp, allocator)
     return C
 end
 
diff --git a/test/gpu.jl b/test/gpu.jl
index 8b716906..81adb21f 100644
--- a/test/gpu.jl
+++ b/test/gpu.jl
@@ -54,6 +54,32 @@ bufferallocator(::Type{ROCArray}; kwargs...) = TensorOperations.AMDBufferAllocat
             tensoradd!(c, a, (p, ()), false, α, β, backend)
         end
 
+        T <: Real || @test compare(AT, C, A) do c, a
+            tensoradd!(c, a, (p, ()), true, α, β, backend)
+        end
+    end
+    # test Diagonal special case
+    sz = (8, 8)
+    p = (2, 1)
+    diag_backends = [BaseCopy(), BaseView()]
+    for backend in diag_backends, T in (Float32, ComplexF32)
+        A = Diagonal(randn(T, sz[1]))
+        C = randn(T, TupleTools.getindices(sz, p))
+
+        @test compare(AT, C, A) do c, a
+            tensoradd!(c, a, (p, ()), false, One(), Zero(), backend)
+        end
+
+        α = rand(T)
+        @test compare(AT, C, A) do c, a
+            tensoradd!(c, a, (p, ()), false, α, Zero(), backend)
+        end
+
+        β = rand(T)
+        @test compare(AT, C, A) do c, a
+            tensoradd!(c, a, (p, ()), false, α, β, backend)
+        end
+
         T <: Real || @test compare(AT, C, A) do c, a
             tensoradd!(c, a, (p, ()), true, α, β, backend)
         end