Adapt extension (#45)

* add Adapt extension

* add adapt tests

* add missing constructor

Author: lkdvos

Project.toml

@@ -15,7 +15,14 @@ TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 
+[weakdeps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+
+[extensions]
+BlockTensorKitAdaptExt = "Adapt"
+
 [compat]
+Adapt = "4"
 Aqua = "0.8"
 BlockArrays = "1"
 Combinatorics = "1"
@@ -34,6 +41,7 @@ VectorInterface = "0.4.8, 0.5"
 julia = "1.10"
 
 [extras]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -42,4 +50,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a"
 
 [targets]
-test = ["Test", "TestExtras", "Random", "Combinatorics", "SafeTestsets", "Aqua"]
+test = ["Test", "TestExtras", "Random", "Combinatorics", "SafeTestsets", "Aqua", "Adapt"]

ext/BlockTensorKitAdaptExt.jl

@@ -0,0 +1,17 @@
+module BlockTensorKitAdaptExt
+
+using TensorKit
+using BlockTensorKit
+using Adapt
+
+function Adapt.adapt_structure(to, x::BlockTensorMap)
+    data′ = map(adapt(to), x.data)
+    return BlockTensorMap(data′, space(x))
+end
+
+function Adapt.adapt_structure(to, x::SparseBlockTensorMap)
+    data′ = Dict(I => adapt(to, v) for (I, v) in x.data)
+    return SparseBlockTensorMap(data′, space(x))
+end
+
+end

src/tensors/sparseblocktensor.jl

@@ -46,6 +46,11 @@ function SparseBlockTensorMap{TT}(
     ) where {E, S, N₁, N₂, N, TT <: AbstractTensorMap{E, S, N₁, N₂}}
     return SparseBlockTensorMap{TT, E, S, N₁, N₂, N}(data, space)
 end
+function SparseBlockTensorMap(
+        data::Dict{CartesianIndex{N}, TT}, space::TensorMapSumSpace{S, N₁, N₂}
+    ) where {E, S, N₁, N₂, N, TT <: AbstractTensorMap{E, S, N₁, N₂}}
+    return SparseBlockTensorMap{TT}(data, space)
+end
 
 function sparseblocktensormaptype(::Type{S}, N₁::Int, N₂::Int, ::Type{T}) where {S, T}
     TT = tensormaptype(S, N₁, N₂, T)

test/abstracttensor/blocktensor.jl

@@ -4,6 +4,7 @@ using TensorKit
 using BlockTensorKit
 using Random
 using Combinatorics
+using Adapt
 
 Vtr = (
     SumSpace(ℂ^3),
@@ -84,6 +85,19 @@ end
     end
 end
 
+@testset "Adapt" begin
+    W = V1 ⊗ V2 ⊗ V3 ← V4 ⊗ V5
+    t1 = rand(Float32, W)
+    for T in (Float64, ComplexF64)
+        t2 = @testinferred adapt(Vector{T}, t1)
+        @test t2 isa BlockTensorMap
+        @test scalartype(t2) == T
+        @test storagetype(t2) == Vector{T}
+        @test space(t1) == space(t2)
+        @test norm(t1) ≈ norm(t2)
+    end
+end
+
 @testset "Basic linear algebra" begin
     W = V1 ⊗ V2 ⊗ V3 ← V4 ⊗ V5
     for T in (Float32, ComplexF64)

test/abstracttensor/sparseblocktensor.jl

@@ -3,6 +3,7 @@ using TensorKit
 using BlockTensorKit
 using Random
 using Combinatorics
+using Adapt
 
 Vtr = (
     SumSpace(ℂ^3),
@@ -90,6 +91,20 @@ end
     end
 end
 
+@testset "Adapt" begin
+    W = V1 ⊗ V2 ⊗ V3 ← V4 ⊗ V5
+    t1 = sprand(Float32, W, 0.5)
+    for T in (Float64, ComplexF64)
+        t2 = @testinferred adapt(Vector{T}, t1)
+        @test t2 isa SparseBlockTensorMap
+        @test scalartype(t2) == T
+        @test storagetype(t2) == Vector{T}
+        @test space(t1) == space(t2)
+        @test norm(t1) ≈ norm(t2)
+    end
+end
+
+
 @testset "Basic linear algebra" begin
     W = V1 ⊗ V2 ⊗ V3 ← V4 ⊗ V5
     for T in (Float32, ComplexF64)