From bcf93b39ee45522e2d04883d536981f427cb6fd3 Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <kslimes@gmail.com>
Date: Thu, 24 Sep 2026 11:53:28 +0200
Subject: [PATCH] Add a rule and tests for TensorMap(d::DiagonalTensorMap)

---
 ext/TensorKitChainRulesCoreExt/constructors.jl | 7 +++++++
 test/chainrules/linalg.jl                      | 3 +++
 2 files changed, 10 insertions(+)

diff --git a/ext/TensorKitChainRulesCoreExt/constructors.jl b/ext/TensorKitChainRulesCoreExt/constructors.jl
index 81cd2dbe9..b6ef760da 100644
--- a/ext/TensorKitChainRulesCoreExt/constructors.jl
+++ b/ext/TensorKitChainRulesCoreExt/constructors.jl
@@ -57,6 +57,13 @@ function ChainRulesCore.rrule(::Type{DiagonalTensorMap}, t::AbstractTensorMap)
     return d, DiagonalTensorMap_pullback
 end
 
+function ChainRulesCore.rrule(::Type{TensorMap}, d::DiagonalTensorMap)
+    t = TensorMap(d)
+    P = ProjectTo(d)
+    TensorMap_pullback(Δt) = NoTangent(), P(unthunk(Δt))
+    return t, TensorMap_pullback
+end
+
 function ChainRulesCore.rrule(::typeof(Base.getproperty), t::TensorMap, prop::Symbol)
     if prop === :data
         function getdata_pullback(Δdata)
diff --git a/test/chainrules/linalg.jl b/test/chainrules/linalg.jl
index cd3da8b2a..880412a3c 100644
--- a/test/chainrules/linalg.jl
+++ b/test/chainrules/linalg.jl
@@ -89,6 +89,9 @@ for V in spacelist
 
                 test_rrule(DiagonalTensorMap, rand!(T1))
                 test_rrule(DiagonalTensorMap, randn!(T))
+
+                test_rrule(TensorMap, D1)
+                test_rrule(TensorMap, D)
             end
         end