Define AD rules for TensorNetwork constructor

ext/TenetChainRulesCoreExt.jl

@@ -28,4 +28,27 @@ ChainRulesCore.rrule(T::Type{<:Tensor}, data, inds; meta...) = T(data, inds; met
 @non_differentiable intersect(s::Base.AbstractVecOrTuple{Symbol}, itrs::Base.AbstractVecOrTuple{Symbol}...)
 @non_differentiable symdiff(s::Base.AbstractVecOrTuple{Symbol}, itrs::Base.AbstractVecOrTuple{Symbol}...)
 
+function ChainRulesCore.ProjectTo(tn::T) where {T<:TensorNetwork}
+    ProjectTo{T}(; tensors = ProjectTo(tn.tensors), metadata = tn.metadata)
+end
+
+function (projector::ProjectTo{T})(dx::Union{T,Tangent{T}}) where {A<:Ansatz,T<:TensorNetwork{A}}
+    TensorNetwork{A}(projector.tensors(dx.tensors); projector.metadata...)
+end
+
+function Base.:+(x::TensorNetwork{A}, Δ::Tangent{TensorNetwork}) where {A<:Ansatz}
+    tensors = map(+, x.tensors, Δ.tensors)
+    TensorNetwork{A}(tensors; x.metadata...)
+end
+
+function ChainRulesCore.frule((_, Δ), T::Type{<:TensorNetwork}, tensors; metadata...)
+    T(tensors; metadata...), Tangent{TensorNetwork}(tensors = Δ)
+end
+
+TensorNetwork_pullback(Δ::Tangent{TensorNetwork}) = (NoTangent(), Δ.tensors)
+TensorNetwork_pullback(Δ::AbstractThunk) = TensorNetwork_pullback(unthunk(Δ))
+function ChainRulesCore.rrule(T::Type{<:TensorNetwork}, tensors; metadata...)
+    T(tensors; metadata...), TensorNetwork_pullback
+end
+
 end

test/integration/ChainRules_test.jl

@@ -12,4 +12,15 @@
         test_frule(Tensor, fill(1.0, 2, 3), Symbol[:i, :j])
         test_rrule(Tensor, fill(1.0, 2, 3), Symbol[:i, :j])
     end
+
+    # NOTE fixes some problems on testing, not sure why
+    Base.collect(tn::TensorNetwork) = tensors(tn)
+
+    @testset "TensorNetwork" begin
+        a = Tensor(rand(4, 2), (:i, :j))
+        b = Tensor(rand(2, 3), (:j, :k))
+
+        test_frule(TensorNetwork, Tensor[a, b])
+        test_rrule(TensorNetwork, Tensor[a, b])
+    end
 end