[NDTensors] UnallocatedArrays and UnspecifiedTypes

NDTensors/Project.toml

@@ -10,6 +10,7 @@ ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 FLoops = "cc61a311-1640-44b5-9fba-1b764f453329"
 Folds = "41a02a25-b8f0-4f67-bc48-60067656b558"

NDTensors/src/NDTensors.jl

@@ -23,6 +23,7 @@ for lib in [
   :AlgorithmSelection,
   :AllocateData,
   :BaseExtensions,
+  :UnspecifiedTypes,
   :SetParameters,
   :BroadcastMapConversion,
   :Unwrap,
@@ -37,6 +38,7 @@ for lib in [
   :SmallVectors,
   :SortedSets,
   :TagSets,
+  :UnallocatedArrays,
 ]
   include("lib/$(lib)/src/$(lib).jl")
   @eval using .$lib: $lib

NDTensors/src/lib/SetParameters/Project.toml

@@ -0,0 +1,8 @@
+[deps]
+PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
+
+[extensions]
+SetParametersFillArraysExt = "FillArrays"
+
+[compat]
+PackageExtensionCompat = "1"

NDTensors/src/lib/SetParameters/TODO.md

@@ -27,3 +27,12 @@ default_parameter(::Type{<:AbstractArray}, ::Position{2}) = 1
 
 nparameters(::Type{<:AbstractArray}) = Val(2)
 ```
+
+https://github.com/ITensor/ITensors.jl/pull/1213/files#r1431708585
+
+# Create generic set_eltype and set_ndims functions which can be defined on 
+# Array structures and use the `set_parameter` 
+```julia
+set_eltype(T::Type, elt::Type) = Error("set_eltype is not defined for datatype $T")
+set_eltype(T::Type{<:Array}, elt::Type) = set_parameter(T, Position{1}(), elt)
+```

NDTensors/src/lib/SetParameters/ext/SetParametersFillArraysExt/SetParametersFillArraysExt.jl

@@ -0,0 +1,4 @@
+module SetParametersFillArraysExt
+include("set_types.jl")
+
+end

NDTensors/src/lib/SetParameters/ext/SetParametersFillArraysExt/set_types.jl

@@ -0,0 +1,77 @@
+using FillArrays: AbstractFill, Fill, Zeros
+using NDTensors.SetParameters: SetParameters, Position
+using NDTensors.UnspecifiedTypes: UnspecifiedZero
+
+# `SetParameters.jl` overloads.
+SetParameters.get_parameter(::Type{<:AbstractFill{P1}}, ::Position{1}) where {P1} = P1
+SetParameters.get_parameter(::Type{<:AbstractFill{<:Any,P2}}, ::Position{2}) where {P2} = P2
+function SetParameters.get_parameter(
+  ::Type{<:AbstractFill{<:Any,<:Any,P3}}, ::Position{3}
+) where {P3}
+  return P3
+end
+
+## Setting paramaters
+# Set parameter 1
+function SetParameters.set_parameter(T::Type{<:AbstractFill}, ::Position{1}, P1)
+  return unspecify_parameters(T){P1}
+end
+function SetParameters.set_parameter(
+  T::Type{<:AbstractFill{<:Any,P2}}, ::Position{1}, P1
+) where {P2}
+  return unspecify_parameters(T){P1,P2}
+end
+function SetParameters.set_parameter(
+  T::Type{<:AbstractFill{<:Any,P2,P3}}, ::Position{1}, P1
+) where {P2,P3}
+  return unspecify_parameters(T){P1,P2,P3}
+end
+
+# Set parameter 2
+function SetParameters.set_parameter(
+  T::Type{<:AbstractFill{P1}}, ::Position{2}, P2
+) where {P1}
+  return unspecify_parameters(T){P1,P2}
+end
+function SetParameters.set_parameter(
+  T::Type{<:AbstractFill{P1,P}}, ::Position{2}, P2
+) where {P1,P}
+  return unspecify_parameters(T){P1,P2}
+end
+function SetParameters.set_parameter(
+  T::Type{<:AbstractFill{P1,P,P3}}, ::Position{2}, P2
+) where {P1,P,P3}
+  return unspecify_parameters(T){P1,P2,P3}
+end
+
+# Set parameter 3
+function SetParameters.set_parameter(
+  T::Type{<:AbstractFill{P1}}, ::Position{3}, P3
+) where {P1}
+  return unspecify_parameters(T){P1,<:Any,P3}
+end
+function SetParameters.set_parameter(
+  T::Type{<:AbstractFill{P1,P2}}, ::Position{3}, P3
+) where {P1,P2}
+  return unspecify_parameters(T){P1,P2,P3}
+end
+function SetParameters.set_parameter(
+  T::Type{<:AbstractFill{P1,P2,P}}, ::Position{3}, P3
+) where {P1,P2,P}
+  return unspecify_parameters(T){P1,P2,P3}
+end
+
+## default parameters
+function SetParameters.default_parameter(::Type{<:AbstractFill}, ::Position{1})
+  return UnspecifiedZero
+end
+SetParameters.default_parameter(::Type{<:AbstractFill}, ::Position{2}) = 0
+SetParameters.default_parameter(::Type{<:AbstractFill}, ::Position{3}) = Tuple{}
+
+SetParameters.nparameters(::Type{<:AbstractFill}) = Val(3)
+
+## These helper functions take a UnallocatedArray type and 
+## remove all the parameters, this way all parameters can be set
+## at once in the `set_parameter` functions above.
+unspecify_parameters(::Type{<:Fill}) = Fill
+unspecify_parameters(::Type{<:Zeros}) = Zeros

NDTensors/src/lib/SetParameters/src/SetParameters.jl

@@ -12,6 +12,15 @@ include("Base/val.jl")
 include("Base/array.jl")
 include("Base/subarray.jl")
 
+## TODO when this is a full package utilize this function to 
+# # enable extensions
+# using PackageExtensionCompat
+# function __init__()
+#   @require_extensions
+# end
+
+include("../ext/SetParametersFillArraysExt/SetParametersFillArraysExt.jl")
+
 export DefaultParameter,
   DefaultParameters,
   Position,

NDTensors/src/lib/TensorAlgebra/test/runtests.jl

@@ -163,9 +163,10 @@ end
       a_dest_tensoroperations = TensorOperations.tensorcontract(
         labels_dest, a1, labels1, a2, labels2
       )
-      @test a_dest ≈ α * a_dest_tensoroperations + β * a_dest_init rtol = default_rtol(
-        elt_dest
-      )
+      ## Here we loosened the tolerance because of some floating point roundoff issue.
+      ## with Float32 numbers
+      @test a_dest ≈ α * a_dest_tensoroperations + β * a_dest_init rtol =
+        10 * default_rtol(elt_dest)
     end
   end
   @testset "qr (eltype=$elt)" for elt in elts

NDTensors/src/lib/UnallocatedArrays/README.md

@@ -0,0 +1,4 @@
+# UnallocatedArrays
+
+A module defining a set of unallocated immutable lazy arrays which will be used to quickly construct 
+tensors and allocating as little data as possible.

NDTensors/src/lib/UnallocatedArrays/src/UnallocatedArrays.jl

@@ -0,0 +1,10 @@
+module UnallocatedArrays
+include("abstractfill/abstractfill.jl")
+
+include("unallocatedfill.jl")
+include("unallocatedzeros.jl")
+include("abstractunallocatedarray.jl")
+include("set_types.jl")
+
+export UnallocatedFill, UnallocatedZeros, alloctype, set_alloctype, allocate
+end

NDTensors/src/lib/UnallocatedArrays/src/abstractfill/abstractfill.jl

@@ -0,0 +1,17 @@
+using FillArrays: AbstractFill
+using NDTensors.SetParameters: Position, get_parameter, set_parameters
+## Here are functions specifically defined for UnallocatedArrays
+## not implemented by FillArrays
+## TODO this might need a more generic name maybe like compute unit
+function alloctype(A::AbstractFill)
+  return A.alloc
+end
+
+## TODO this fails if the parameter is a type
+function alloctype(Atype::Type{<:AbstractFill})
+  return get_parameter(Atype, Position{4}())
+end
+
+set_eltype(T::Type{<:AbstractFill}, elt::Type) = set_parameters(T, Position{1}(), elt)
+set_ndims(T::Type{<:AbstractFill}, n) = set_parameters(T, Position{2}(), n)
+set_axestype(T::Type{<:AbstractFill}, ax::Type) = set_parameters(T, Position{3}(), ax)

NDTensors/src/lib/UnallocatedArrays/src/abstractunallocatedarray.jl

@@ -0,0 +1,64 @@
+using FillArrays: FillArrays, getindex_value
+using NDTensors.SetParameters: set_parameters
+using Adapt: adapt
+
+const UnallocatedArray{ElT,N,AxesT,AllocT} = Union{
+  UnallocatedFill{ElT,N,AxesT,AllocT},UnallocatedZeros{ElT,N,AxesT,AllocT}
+}
+
+@inline Base.axes(A::UnallocatedArray) = axes(parent(A))
+Base.size(A::UnallocatedArray) = size(parent(A))
+function FillArrays.getindex_value(A::UnallocatedArray)
+  return getindex_value(parent(A))
+end
+
+function Base.complex(A::UnallocatedArray)
+  return complex(eltype(A)).(A)
+end
+
+function Base.transpose(a::UnallocatedArray)
+  return set_alloctype(transpose(parent(a)), alloctype(a))
+end
+
+function Base.adjoint(a::UnallocatedArray)
+  return set_alloctype(adjoint(parent(a)), alloctype(a))
+end
+
+function set_alloctype(T::Type{<:UnallocatedArray}, alloc::Type{<:AbstractArray})
+  return set_parameters(T, Position{4}(), alloc)
+end
+
+## This overloads the definition defined in `FillArrays.jl`
+for STYPE in (:AbstractArray, :AbstractFill)
+  @eval begin
+    @inline $STYPE{T}(F::UnallocatedArray{T}) where {T} = F
+    @inline $STYPE{T,N}(F::UnallocatedArray{T,N}) where {T,N} = F
+  end
+end
+
+function allocate(f::UnallocatedArray)
+  a = similar(f)
+  fill!(a, getindex_value(f))
+  return a
+end
+
+function allocate(arraytype::Type{<:AbstractArray}, elt::Type, axes)
+  ## TODO rewrite this using set_eltype and set_ndims functions
+  ## currently these functions are defined in `NDTensors`
+  ## In the future they should be defined in `SetParameters`
+  ArrayT = set_parameters(arraytype, Position{1}(), elt)
+  return similar(ArrayT, axes)
+end
+
+function Base.similar(f::UnallocatedArray, elt::Type, axes::Tuple{Int64,Vararg{Int64}})
+  return allocate(alloctype(f), elt, axes)
+end
+
+## TODO fix this because reshape loses alloctype
+#FillArrays.reshape(a::Union{<:UnallocatedFill, <:UnallocatedZeros}, dims) = set_alloctype(reshape(parent(a), dims), allocate(a))
+
+# function Adapt.adapt_storage(to::Type{<:AbstractArray}, x::Union{<:UnallocatedFill, <:UnallocatedZeros})
+#   return set_alloctype(parent(x), to)
+# end
+
+# function Adapt.adapt_storage(to::Type{<:Number}, x::)

NDTensors/src/lib/UnallocatedArrays/src/set_types.jl

@@ -0,0 +1,45 @@
+using NDTensors.SetParameters: SetParameters, Position
+using NDTensors.UnspecifiedTypes: UnspecifiedArray, UnspecifiedNumber, UnspecifiedZero
+# ## TODO make unit tests for all of these functions
+## TODO All I need to do is overload AbstractFill functions with 4 parameters
+# `SetParameters.jl` overloads.
+function SetParameters.get_parameter(
+  ::Type{<:UnallocatedArray{<:Any,<:Any,<:Any,P4}}, ::Position{4}
+) where {P4}
+  return P4
+end
+
+# ## Setting paramaters
+function SetParameters.set_parameter(
+  T::Type{<:UnallocatedArray{P,P2,P3,P4}}, ::Position{1}, P1
+) where {P,P2,P3,P4}
+  return unspecify_parameters(T){P1,P2,P3,P4}
+end
+
+function SetParameters.set_parameter(
+  T::Type{<:UnallocatedArray{P1,P,P3,P4}}, ::Position{2}, P2
+) where {P1,P,P3,P4}
+  return unspecify_parameters(T){P1,P2,P3,P4}
+end
+
+function SetParameters.set_parameter(
+  T::Type{<:UnallocatedArray{P1,P2,P,P4}}, ::Position{3}, P3
+) where {P1,P2,P,P4}
+  return unspecify_parameters(T){P1,P2,P3,P4}
+end
+
+function SetParameters.set_parameter(
+  T::Type{<:UnallocatedArray{P1,P2,P3}}, ::Position{4}, P4
+) where {P1,P2,P3}
+  return unspecify_parameters(T){P1,P2,P3,P4}
+end
+
+# ## default parameters
+function SetParameters.default_parameter(::Type{<:UnallocatedArray}, ::Position{4})
+  return UnspecifiedArray{UnspecifiedNumber{UnspecifiedZero},0}
+end
+
+SetParameters.nparameters(::Type{<:UnallocatedArray}) = Val(4)
+
+unspecify_parameters(::Type{<:UnallocatedFill}) = UnallocatedFill
+unspecify_parameters(::Type{<:UnallocatedZeros}) = UnallocatedZeros

NDTensors/src/lib/UnallocatedArrays/src/unallocatedfill.jl

@@ -0,0 +1,67 @@
+using FillArrays: FillArrays, AbstractFill, Fill, broadcasted_fill, kron_fill, mult_fill
+using NDTensors.SetParameters: Position, set_parameters
+
+struct UnallocatedFill{ElT,N,Axes,Alloc} <: AbstractFill{ElT,N,Axes}
+  f::Fill{ElT,N,Axes}
+  alloc::Alloc
+end
+
+function UnallocatedFill{ElT,N,Axes}(f::Fill, alloc::Type) where {ElT,N,Axes}
+  return new{ElT,N,Axes,Type{alloc}}(f, alloc)
+end
+
+function UnallocatedFill{ElT,N}(f::Fill, alloc) where {ElT,N}
+  return UnallocatedFill{ElT,N,typeof(axes(f))}(f, alloc)
+end
+
+function UnallocatedFill{ElT}(f::Fill, alloc) where {ElT}
+  return UnallocatedFill{ElT,ndims(f)}(f, alloc)
+end
+
+set_alloctype(f::Fill, alloc::Type) = UnallocatedFill(f, alloc)
+
+Base.parent(F::UnallocatedFill) = F.f
+
+Base.convert(::Type{<:UnallocatedFill}, A::UnallocatedFill) = A
+
+#############################################
+# Arithmatic
+
+# mult_fill(a, b, val, ax) = Fill(val, ax)
+function FillArrays.mult_fill(a::UnallocatedFill, b, val, ax)
+  return UnallocatedFill(Fill(val, ax), alloctype(a))
+end
+FillArrays.mult_fill(a, b::UnallocatedFill, val, ax) = mult_fill(b, a, val, ax)
+function FillArrays.mult_fill(a::UnallocatedFill, b::UnallocatedFill, val, ax)
+  @assert alloctype(a) == alloctype(b)
+  return UnallocatedFill(Fill(val, ax), alloctype(a))
+end
+
+function FillArrays.broadcasted_fill(f, a::UnallocatedFill, val, ax)
+  return UnallocatedFill(Fill(val, ax), alloctype(a))
+end
+function FillArrays.broadcasted_fill(f, a::UnallocatedFill, b::UnallocatedFill, val, ax)
+  @assert alloctype(a) == alloctype(b)
+  return UnallocatedFill(Fill(val, ax), alloctype(a))
+end
+
+function FillArrays.broadcasted_fill(f, a::UnallocatedFill, b, val, ax)
+  return UnallocatedFill(Fill(val, ax), alloctype(a))
+end
+function FillArrays.broadcasted_fill(f, a, b::UnallocatedFill, val, ax)
+  return broadcasted_fill(f, b, a, val, ax)
+end
+
+function FillArrays.kron_fill(a::UnallocatedFill, b::UnallocatedFill, val, ax)
+  @assert alloctype(a) == alloctype(b)
+  return UnallocatedFill(Fill(val, ax), alloctype(a))
+end
+
+Base.:+(A::UnallocatedFill, B::UnallocatedFill) = A .+ B
+
+function Base.Broadcast.broadcasted(
+  ::Base.Broadcast.DefaultArrayStyle, op, r::UnallocatedFill
+)
+  f = op.(parent(r))
+  return set_alloctype(f, set_parameters(alloctype(r), Position{1}(), eltype(f)))
+end