Add a disamgiguating conversion (#47)

* Add a disamgiguating conversion

* Add comment

* Incremental progress on getting rid of allowscalar

* Support BlockArrays in KernelAbstractions bcasting with new extension

* Add test for different eltype

* Update Project.toml

* Update src/tensors/tensoroperations.jl

Co-authored-by: Lukas Devos <ldevos98@gmail.com>

* Move element type test out

---------

Co-authored-by: Lukas Devos <ldevos98@gmail.com>

Author: kshyatt

Project.toml

@@ -17,21 +17,25 @@ VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 
 [weakdeps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 
 [extensions]
 BlockTensorKitAdaptExt = "Adapt"
+BlockTensorKitGPUArraysExt = "GPUArrays"
 
 [compat]
 Adapt = "4"
 Aqua = "0.8"
 BlockArrays = "1"
 Combinatorics = "1"
 Compat = "4.13"
+GPUArrays = "11.4.1"
+JLArrays = "0.3"
 LinearAlgebra = "1"
 MatrixAlgebraKit = "0.6"
 Random = "1"
 SafeTestsets = "0.1"
-Strided = "2"
+Strided = "2.3.3"
 TensorKit = "0.16.1"
 TensorOperations = "5"
 Test = "1"
@@ -44,10 +48,12 @@ julia = "1.10"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
+GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+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"]

ext/BlockTensorKitGPUArraysExt.jl

@@ -0,0 +1,11 @@
+module BlockTensorKitGPUArraysExt
+
+using BlockTensorKit, BlockArrays, GPUArrays, Strided
+using Strided: StridedViews
+using GPUArrays: KernelAbstractions
+
+function KernelAbstractions.get_backend(BA::BlockArrays.BlockArray{T, N, A}) where {T, N, A <: AbstractArray{<:StridedView{T, N, <:AnyGPUArray}}}
+    return KernelAbstractions.get_backend(first(BA.blocks))
+end
+
+end

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/blocktensor.jl

@@ -114,7 +114,7 @@ function BlockTensorMap(t::AbstractTensorMap, space::TensorMapSumSpace)
     TT = tensormaptype(spacetype(t), numout(t), numin(t), storagetype(t))
     tdst = BlockTensorMap{TT}(undef, space)
     for (f₁, f₂) in fusiontrees(tdst)
-        tdst[f₁, f₂] .= t[f₁, f₂]
+        copy!(tdst[f₁, f₂], t[f₁, f₂])
     end
     return tdst
 end

src/tensors/tensoroperations.jl

@@ -15,6 +15,10 @@ function TO.tensoradd_type(TC, A::AdjointBlockTensorMap, pA::Index2Tuple, conjA:
     return TO.tensoradd_type(TC, A', adjointtensorindices(A, pA), !conjA)
 end
 
+function TO.tensorscalar(t::AbstractBlockTensorMap{T, S, 0, 0}) where {T, S}
+    return nonzero_length(t) == 0 ? zero(T) : TO.tensorscalar(only(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,7 +83,20 @@ 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}}
+        jl_bt1′ = @constinferred convert(TT, jl_bt1)
+        jl_bt1″ = @inferred BlockTensorMap(jl_bt1′, W)
+        @test jl_bt1 ≈ jl_bt1″
     end
+    # test conversion to TensorMap with a different element type
+    t1 = rand(ComplexF32, W)
+    TT = TensorKit.TensorMap{ComplexF64, spacetype(t1), numout(t1), numin(t1), Vector{ComplexF64}}
+    t1′ = @constinferred convert(TT, t1)
+    @test norm(t1) ≈ norm(t1′)
+    t1″ = @inferred BlockTensorMap(t1′, W)
+    @test t1 ≈ t1″
 end
 
 @testset "Adapt" begin