Add support for Adapt

Author: mtfishman

Project.toml

@@ -1,15 +1,23 @@
 name = "SerializedArrays"
 uuid = "621c0da3-e96e-4f80-bd06-5ae31cdfcb39"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
-version = "0.1.0"
+version = "0.1.1"
 
 [deps]
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 DiskArrays = "3c3547ce-8d99-4f5e-a174-61eb10b00ae3"
-LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 
+[weakdeps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+
+[extensions]
+SerializedArraysAdaptExt = "Adapt"
+SerializedArraysLinearAlgebraExt = "LinearAlgebra"
+
 [compat]
+Adapt = "4.3.0"
 ConstructionBase = "1.5.8"
 DiskArrays = "0.4.12"
 LinearAlgebra = "1.10"

ext/SerializedArraysAdaptExt/SerializedArraysAdaptExt.jl

@@ -0,0 +1,10 @@
+module SerializedArraysAdaptExt
+
+using Adapt: Adapt
+using SerializedArrays: SerializedArray
+
+function Adapt.adapt_storage(arrayt::Type{<:SerializedArray}, a::AbstractArray)
+  return convert(arrayt, a)
+end
+
+end

ext/SerializedArraysLinearAlgebraExt/SerializedArraysLinearAlgebraExt.jl

@@ -0,0 +1,13 @@
+module SerializedArraysLinearAlgebraExt
+
+using LinearAlgebra: LinearAlgebra, mul!
+using SerializedArrays: SerializedArray
+
+function LinearAlgebra.mul!(
+  a_dest::AbstractMatrix, a1::SerializedArray, a2::SerializedArray, α::Number, β::Number
+)
+  mul!(a_dest, copy(a1), copy(a2), α, β)
+  return a_dest
+end
+
+end

src/SerializedArrays.jl

@@ -2,7 +2,6 @@ module SerializedArrays
 
 using ConstructionBase: constructorof
 using DiskArrays: DiskArrays, AbstractDiskArray, Unchunked
-using LinearAlgebra: LinearAlgebra, mul!
 using Serialization: deserialize, serialize
 
 struct SerializedArray{T,N,A<:AbstractArray{T,N},Axes} <: AbstractDiskArray{T,N}
@@ -14,21 +13,25 @@ Base.axes(a::SerializedArray) = getfield(a, :axes)
 arraytype(a::SerializedArray{<:Any,<:Any,A}) where {A} = A
 
 function SerializedArray(file::String, a::AbstractArray)
-  serialize(file, a)
+  serialize(file, vec(a))
   ax = axes(a)
   return SerializedArray{eltype(a),ndims(a),typeof(a),typeof(ax)}(file, ax)
 end
 function SerializedArray(a::AbstractArray)
   return SerializedArray(tempname(), a)
 end
 
+function Base.convert(arrayt::Type{<:SerializedArray}, a::AbstractArray)
+  return arrayt(a)
+end
+
 function Base.similar(a::SerializedArray, elt::Type, dims::Tuple{Vararg{Int}})
   return constructorof(arraytype(a)){elt}(undef, dims...)
 end
 
 function Base.copy(a::SerializedArray)
   arrayt = arraytype(a)
-  return convert(arrayt, deserialize(file(a)))::arrayt
+  return convert(arrayt, reshape(deserialize(file(a)), axes(a)))::arrayt
 end
 
 Base.size(a::SerializedArray) = length.(axes(a))
@@ -106,15 +109,4 @@ function Base.copy(broadcasted::Broadcasted{SerializedArrayStyle{N}}) where {N}
   return BroadcastSerializedArray(flatten(broadcasted))
 end
 
-#
-# LinearAlgebra
-#
-
-function LinearAlgebra.mul!(
-  a_dest::AbstractMatrix, a1::SerializedArray, a2::SerializedArray, α::Number, β::Number
-)
-  mul!(a_dest, copy(a1), copy(a2), α, β)
-  return a_dest
-end
-
 end