[ITensorMPS] Code reorganization (#1406)

Author: mtfishman

Project.toml

@@ -29,19 +29,22 @@ Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 Zeros = "bd1ec220-6eb4-527a-9b49-e79c3db6233b"
-ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [weakdeps]
+# ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 Observers = "338f10d5-c7f1-4033-a7d1-f9dec39bcaa0"
 PackageCompiler = "9b87118b-4619-50d2-8e1e-99f35a4d4d9d"
 VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
+ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [extensions]
+# ITensorsChainRulesCoreExt = "ChainRulesCore"
 ITensorsHDF5Ext = "HDF5"
 ITensorsObserversExt = "Observers"
 ITensorsPackageCompilerExt = "PackageCompiler"
 ITensorsVectorInterfaceExt = "VectorInterface"
+ITensorsZygoteRulesExt = "ZygoteRules"
 
 [compat]
 Adapt = "3.5, 4"
@@ -76,7 +79,9 @@ ZygoteRules = "0.2.2"
 julia = "1.6"
 
 [extras]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 Observers = "338f10d5-c7f1-4033-a7d1-f9dec39bcaa0"
 PackageCompiler = "9b87118b-4619-50d2-8e1e-99f35a4d4d9d"
 VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
+ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"

ext/ITensorsChainRulesCoreExt/ITensorsChainRulesCoreExt.jl

@@ -0,0 +1,20 @@
+module ITensorsChainRulesCoreExt
+using ChainRulesCore
+import ChainRulesCore: rrule
+using ITensors
+using ITensors: Indices
+using ITensors.Adapt
+using ITensors.NDTensors
+using ITensors.NDTensors: datatype
+using ITensors.Ops
+include("utils.jl")
+include("projection.jl")
+include("NDTensors/tensor.jl")
+include("NDTensors/dense.jl")
+include("indexset.jl")
+include("itensor.jl")
+include("LazyApply/LazyApply.jl")
+include("non_differentiable.jl")
+include("itensormps.jl")
+include("smallstrings.jl")
+end

src/lib/ITensorChainRules/src/LazyApply/LazyApply.jl renamed to ext/ITensorsChainRulesCoreExt/LazyApply/LazyApply.jl

@@ -0,0 +1,32 @@
+for fname in (
+  :prime,
+  :setprime,
+  :noprime,
+  :replaceprime,
+  :swapprime,
+  :addtags,
+  :removetags,
+  :replacetags,
+  :settags,
+  :swaptags,
+  :replaceind,
+  :replaceinds,
+  :swapind,
+  :swapinds,
+)
+  @eval begin
+    function rrule(f::typeof($fname), x::ITensor, a...; kwargs...)
+      y = f(x, a...; kwargs...)
+      function f_pullback(ȳ)
+        x̄ = replaceinds(unthunk(ȳ), inds(y) => inds(x))
+        ā = map_notangent(a)
+        return (NoTangent(), x̄, ā...)
+      end
+      return y, f_pullback
+    end
+  end
+end
+
+rrule(::typeof(adjoint), x::ITensor) = rrule(prime, x)
+
+@non_differentiable permute(::Indices, ::Indices)

src/lib/ITensorChainRules/src/NDTensors/dense.jl renamed to ext/ITensorsChainRulesCoreExt/NDTensors/dense.jl

@@ -0,0 +1,12 @@
+using ITensors: ITensors
+include(
+  joinpath(
+    pkgdir(ITensors),
+    "src",
+    "lib",
+    "ITensorMPS",
+    "ext",
+    "ITensorMPSChainRulesCoreExt",
+    "ITensorMPSChainRulesCoreExt.jl",
+  ),
+)

src/lib/ITensorChainRules/src/NDTensors/tensor.jl renamed to ext/ITensorsChainRulesCoreExt/NDTensors/tensor.jl

@@ -0,0 +1,19 @@
+using ChainRulesCore: @non_differentiable
+using ITensors:
+  ITensors, Index, addtags, commoninds, dag, delta, inds, noncommoninds, onehot, uniqueinds
+using ITensors.TagSets: TagSet
+
+@non_differentiable map_notangent(::Any)
+@non_differentiable Index(::Any...)
+@non_differentiable delta(::Any...)
+@non_differentiable dag(::Index)
+@non_differentiable inds(::Any...)
+@non_differentiable commoninds(::Any...)
+@non_differentiable noncommoninds(::Any...)
+@non_differentiable uniqueinds(::Any...)
+@non_differentiable addtags(::TagSet, ::Any)
+@non_differentiable ITensors.filter_inds_set_function(::Function, ::Function, ::Any...)
+@non_differentiable ITensors.filter_inds_set_function(::Function, ::Any...)
+@non_differentiable ITensors.indpairs(::Any...)
+@non_differentiable onehot(::Any...)
+@non_differentiable Base.convert(::Type{TagSet}, str::String)