Add a disamgiguating conversion

Author: kshyatt

Project.toml

@@ -27,6 +27,7 @@ Aqua = "0.8"
 BlockArrays = "1"
 Combinatorics = "1"
 Compat = "4.13"
+JLArrays = "0.3"
 LinearAlgebra = "1"
 MatrixAlgebraKit = "0.6"
 Random = "1"
@@ -44,10 +45,11 @@ julia = "1.10"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a"
 
 [targets]
-test = ["Test", "TestExtras", "Random", "Combinatorics", "SafeTestsets", "Aqua", "Adapt"]
+test = ["Test", "TestExtras", "Random", "Combinatorics", "SafeTestsets", "Aqua", "Adapt", "JLArrays"]

src/tensors/abstractblocktensor/conversion.jl

@@ -1,37 +1,50 @@
 # Conversion
 # ----------
-function Base.convert(::Type{TensorMap}, t::AbstractBlockTensorMap)
-    S = spacetype(t)
-    N₁, N₂ = numout(t), numin(t)
-    cod = ProductSpace{S, N₁}(oplus.(codomain(t).spaces))
-    dom = ProductSpace{S, N₂}(oplus.(domain(t).spaces))
-    tdst = similar(t, cod ← dom)
-
-    issparse(t) && zerovector!(tdst)
 
+function _copy_subblocks!(tdst, tsrc)
+    S = spacetype(tsrc)
+    N₁, N₂ = numout(tsrc), numin(tsrc)
     for ((f₁, f₂), arr) in subblocks(tdst)
         blockax = ntuple(N₁ + N₂) do i
             return if i <= N₁
-                blockedrange(map(Base.Fix2(dim, f₁.uncoupled[i]), space(t, i)))
+                blockedrange(map(Base.Fix2(dim, f₁.uncoupled[i]), space(tsrc, i)))
             else
-                blockedrange(map(Base.Fix2(dim, f₂.uncoupled[i - N₁]), space(t, i)'))
+                blockedrange(map(Base.Fix2(dim, f₂.uncoupled[i - N₁]), space(tsrc, i)'))
             end
         end
 
-        for (k, v) in nonzero_pairs(t)
+        for (k, v) in nonzero_pairs(tsrc)
             indices = getindex.(blockax, Block.(Tuple(k)))
             arr_slice = arr[indices...]
             # need to check for empty since fusion tree pair might not be present
             isempty(arr_slice) || copy!(arr_slice, v[f₁, f₂])
         end
     end
+    return tdst
+end
 
+function Base.convert(::Type{TensorMap}, t::AbstractBlockTensorMap)
+    S = spacetype(t)
+    N₁, N₂ = numout(t), numin(t)
+    cod = ProductSpace{S, N₁}(oplus.(codomain(t).spaces))
+    dom = ProductSpace{S, N₂}(oplus.(domain(t).spaces))
+    tdst = TensorKit.TensorMapWithStorage{scalartype(t), storagetype(t)}(undef, cod, dom)
+
+    issparse(t) && zerovector!(tdst)
+    _copy_subblocks!(tdst, t)
     return tdst
 end
 
-function Base.convert(::Type{T}, t::AbstractBlockTensorMap) where {T <: TensorMap}
-    tdst = convert(TensorMap, t)
-    return convert(T, tdst)
+function Base.convert(::Type{TT}, t::AbstractBlockTensorMap) where {TT <: TensorKit.TensorMap}
+    S = spacetype(t)
+    N₁, N₂ = numout(t), numin(t)
+    cod = ProductSpace{S, N₁}(oplus.(codomain(t).spaces))
+    dom = ProductSpace{S, N₂}(oplus.(domain(t).spaces))
+    tdst = TT(undef, cod ← dom)
+    issparse(t) && zerovector!(tdst)
+
+    _copy_subblocks!(tdst, t)
+    return tdst
 end
 
 function Base.convert(::Type{TT}, t::AbstractTensorMap) where {TT <: AbstractBlockTensorMap}

src/tensors/tensoroperations.jl

@@ -15,6 +15,13 @@ function TO.tensoradd_type(TC, A::AdjointBlockTensorMap, pA::Index2Tuple, conjA:
     return TO.tensoradd_type(TC, A', adjointtensorindices(A, pA), !conjA)
 end
 
+# copy blocks back to CPU/collect them into an array
+# seems necessary for GPU-backed BlockTensorMaps but
+# maybe not the most efficient approach?
+function TO.tensorscalar(t::AbstractBlockTensorMap{T, S, 0, 0}) where {T, S}
+    return prod(TO.tensorscalar, nonzero_values(t))
+end
+
 # tensoralloc_contract
 # --------------------
 for TTA in (:AbstractTensorMap, :AbstractBlockTensorMap), TTB in (:AbstractTensorMap, :AbstractBlockTensorMap)

test/abstracttensor/blocktensor.jl

@@ -5,6 +5,7 @@ using BlockTensorKit
 using Random
 using Combinatorics
 using Adapt
+using JLArrays
 
 Vtr = (
     SumSpace(ℂ^3),
@@ -82,6 +83,14 @@ end
         t2″ = @inferred BlockTensorMap(t2′, W)
         @test t1 ≈ t1″
         @test t2 ≈ t2″
+        # test conversion to TensorMap that isn't backed by a Vector
+        jl_bt1 = rand(JLVector{T}, W)
+        TT = TensorKit.TensorMap{T, spacetype(t1′), numout(t1′), numin(t1′), JLVector{T}}
+        JLArrays.@allowscalar begin # to avoid scalar indexing error in Strided
+            jl_bt1′ = @constinferred convert(TT, jl_bt1)
+            jl_bt1″ = @inferred BlockTensorMap(jl_bt1′, W)
+        end
+        @test jl_bt1 ≈ jl_bt1″
     end
 end