Fix adapt functions for UnallocatedZeros

ITensor · Sep 20, 2023 · ca75862 · ca75862
1 parent 2f38cef
commit ca75862
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 5 deletions.
diff --git a/NDTensors/ext/NDTensorsCUDAExt/adapt.jl b/NDTensors/ext/NDTensorsCUDAExt/adapt.jl
@@ -18,3 +18,18 @@ function Adapt.adapt_storage(adaptor::NDTensorCuArrayAdaptor, xs::AbstractArray)
   BufT = buffertype(adaptor)
   return isbits(xs) ? xs : CuArray{ElT,1,BufT}(xs)
 end
+
+function Adapt.adapt_storage(arraytype::Type{<:CuArray}, xs::NDTensors.UnallocatedZeros)
+  arraytype_specified_1 = set_unspecified_parameters(
+    arraytype, Position(1), get_parameter(xs, Position(1))
+  )
+  arraytype_specified_2 = set_unspecified_parameters(
+    arraytype_specified_1, Position(2), get_parameter(xs, Position(2))
+  )
+  arraytype_specified_3 = set_unspecified_parameters(
+    arraytype_specified_2, Position(3), get_parameter(xs, Position(3))
+  )
+  elt = get_parameter(arraytype_specified_3, Position(1))
+  N = get_parameter(arraytype_specified_3, Position(2))
+  return NDTensors.UnallocatedZeros{elt, N, CUDA.CuArray{elt, N, default_parameter(CuArray, Position(3))}}(size(xs))
+end
diff --git a/NDTensors/ext/NDTensorsMetalExt/adapt.jl b/NDTensors/ext/NDTensorsMetalExt/adapt.jl
@@ -14,5 +14,22 @@ function Adapt.adapt_storage(arraytype::Type{<:MtlArray}, xs::AbstractArray)
   arraytype_specified_3 = set_unspecified_parameters(
     arraytype_specified_2, Position(3), get_parameter(xs, Position(3))
   )
+  @show convert(arraytype_specified_3, xs)
   return isbitstype(typeof(xs)) ? xs : convert(arraytype_specified_3, xs)
 end
+
+function Adapt.adapt_storage(arraytype::Type{<:MtlArray}, xs::NDTensors.UnallocatedZeros)
+  arraytype_specified_1 = set_unspecified_parameters(
+    arraytype, Position(1), get_parameter(xs, Position(1))
+  )
+  arraytype_specified_2 = set_unspecified_parameters(
+    arraytype_specified_1, Position(2), get_parameter(xs, Position(2))
+  )
+  arraytype_specified_3 = set_unspecified_parameters(
+    arraytype_specified_2, Position(3), get_parameter(xs, Position(3))
+  )
+  elt = get_parameter(arraytype_specified_3, Position(1))
+  N = get_parameter(arraytype_specified_3, Position(2))
+
+  return NDTensors.UnallocatedZeros{elt, N, Metal.MtlArray{elt, N, default_parameter(MtlArray, Position(3))}}(size(xs))
+end
diff --git a/NDTensors/src/zeros/set_types.jl b/NDTensors/src/zeros/set_types.jl
@@ -1,10 +1,10 @@
 import .SetParameters: set_parameter, nparameters, default_parameter
 
 # `SetParameters.jl` overloads.
-get_parameter(::Type{<:UnallocatedZeros{P1}}, ::Position{1}) where {P1} = P1
-get_parameter(::Type{<:UnallocatedZeros{<:Any,P2}}, ::Position{2}) where {P2} = P2
-get_parameter(::Type{<:UnallocatedZeros{<:Any,<:Any,P3}}, ::Position{3}) where {P3} = P3
-function get_parameter(
+NDTensors.SetParameters.get_parameter(::Type{<:UnallocatedZeros{P1}}, ::Position{1}) where {P1} = P1
+NDTensors.SetParameters.get_parameter(::Type{<:UnallocatedZeros{<:Any,P2}}, ::Position{2}) where {P2} = P2
+NDTensors.SetParameters.get_parameter(::Type{<:UnallocatedZeros{<:Any,<:Any,P3}}, ::Position{3}) where {P3} = P3
+function NDTensors.SetParameters.get_parameter(
   ::Type{<:UnallocatedZeros{<:Any,<:Any,<:Any,P4}}, ::Position{4}
 ) where {P4}
   return P4
@@ -36,12 +36,13 @@ end
 function set_parameter(
   ::Type{<:UnallocatedZeros{P1,P2,P3,<:Any}}, ::Position{4}, P4
 ) where {P1,P2,P3}
+@show P4
   return UnallocatedZeros{P1,P2,P3,P4}
 end
 
 default_parameter(::Type{<:UnallocatedZeros}, ::Position{1}) = Float64
 default_parameter(::Type{<:UnallocatedZeros}, ::Position{2}) = 1
 default_parameter(::Type{<:UnallocatedZeros}, ::Position{3}) = Tuple{Base.OneTo{Int64}}
-default_parameter(::Type{<:UnallocatedZeros}, ::Position{4}) = Vector{Float64}
+default_parameter(::Type{<:UnallocatedZeros}, ::Position{4}) = Vector{default_parameter(UnallocatedZeros, Position(1))}
 
 nparameters(::Type{<:UnallocatedZeros}) = Val(4)