Fix Adapt ext

Author: kshyatt

ext/MPSKitAdaptExt.jl

@@ -1,6 +1,6 @@
 module MPSKitAdaptExt
 
-using TensorKit: space, spacetype
+using TensorKit: space, spacetype, similarstoragetype, scalartype
 using MPSKit
 using BlockTensorKit: nonzero_pairs
 using Adapt
@@ -32,7 +32,8 @@ end
 
 Adapt.adapt_structure(to, mpo::MPO) = MPO(map(adapt(to), mpo.O))
 
-function Adapt.adapt_structure(::Type{TorA}, W::MPSKit.JordanMPOTensor) where {TorA <: Union{Number, DenseVector{<:Number}}}
+function Adapt.adapt_structure(to, W::MPSKit.JordanMPOTensorMap)
+    TorA = similarstoragetype(to, scalartype(W))
     TT = MPSKit.jordanmpotensortype(spacetype(W), TorA)
     W′ = TT(undef, space(W))
     ad = adapt(TorA)
@@ -49,26 +50,9 @@ function Adapt.adapt_structure(::Type{TorA}, W::MPSKit.JordanMPOTensor) where {T
     for (k, v) in nonzero_pairs(W.D)
         W′.D[k] = ad(v)
     end
-
     return W′
 end
 Adapt.adapt_structure(to, mpo::MPOHamiltonian) = MPOHamiltonian(map(adapt(to), mpo.W))
-
-function Adapt.adapt_structure(to, x::MPSKit.MPOHamiltonian{TO}) where {TO}
-    terms′ = map(w -> adapt(to, w), x.W)
-    return MPSKit.MPOHamiltonian(terms′)
-end
-
-function Adapt.adapt_structure(to, x::MPSKit.PeriodicArray)
-    return MPSKit.PeriodicArray(map(x_ -> adapt(to, x_), x.data))
-end
-
-function Adapt.adapt_structure(to, x::MPSKit.InfiniteMPS{A, B}) where {A, S, B <: MPSKit.MPSBondTensor{S}}
-    AL′ = adapt(to, x.AL)
-    AR′ = adapt(to, x.AR)
-    AC′ = adapt(to, x.AC)
-    C′ = adapt(to, x.C)
-    return MPSKit.InfiniteMPS(AL′, AR′, C′, AC′)
-end
+Adapt.adapt_structure(to, x::MPSKit.PeriodicArray) = MPSKit.PeriodicArray(map(adapt(to), x.data))
 
 end