diff --git a/NDTensors/Project.toml b/NDTensors/Project.toml index ca244f17b3..a18cec49c5 100644 --- a/NDTensors/Project.toml +++ b/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" FLoops = "cc61a311-1640-44b5-9fba-1b764f453329" Folds = "41a02a25-b8f0-4f67-bc48-60067656b558" Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" diff --git a/NDTensors/src/NDTensors.jl b/NDTensors/src/NDTensors.jl index 052927677f..ad9059d887 100644 --- a/NDTensors/src/NDTensors.jl +++ b/NDTensors/src/NDTensors.jl @@ -37,6 +37,8 @@ for lib in [ :SmallVectors, :SortedSets, :TagSets, + :UnspecifiedTypes, + :UnallocatedArrays, ] include("lib/$(lib)/src/$(lib).jl") @eval using .$lib: $lib diff --git a/NDTensors/src/lib/NamedDimsArrays/ext/NamedDimsArraysTensorAlgebraExt/test/runtests.jl b/NDTensors/src/lib/NamedDimsArrays/ext/NamedDimsArraysTensorAlgebraExt/test/runtests.jl index 31106b56d8..2690f43ccc 100644 --- a/NDTensors/src/lib/NamedDimsArrays/ext/NamedDimsArraysTensorAlgebraExt/test/runtests.jl +++ b/NDTensors/src/lib/NamedDimsArrays/ext/NamedDimsArraysTensorAlgebraExt/test/runtests.jl @@ -3,17 +3,23 @@ using Test: @test, @testset, @test_broken using NDTensors.NamedDimsArrays: named, unname using NDTensors.TensorAlgebra: TensorAlgebra using LinearAlgebra: qr -@testset "NamedDimsArraysTensorAlgebraExt contract (eltype=$(elt))" for elt in ( - Float32, ComplexF32, Float64, ComplexF64 -) + +using NDTensors: NDTensors +using GPUArraysCore: @allowscalar +include(joinpath(pkgdir(NDTensors), "test", "NDTensorsTestUtils", "NDTensorsTestUtils.jl")) +using .NDTensorsTestUtils: devices_list +@testset "NamedDimsArraysTensorAlgebraExt contract (eltype=$(elt))" for dev in + devices_list(ARGS), + elt in (Float32, ComplexF32, Float64, ComplexF64) + i = named(2, "i") j = named(2, "j") k = named(2, "k") - na1 = randn(elt, i, j) - na2 = randn(elt, j, k) + na1 = dev(randn(elt, i, j)) + na2 = dev(randn(elt, j, k)) na_dest = TensorAlgebra.contract(na1, na2) @test eltype(na_dest) === elt - @test unname(na_dest, (i, k)) ≈ unname(na1) * unname(na2) + @allowscalar @test unname(na_dest, (i, k)) ≈ unname(na1) * unname(na2) end @testset "NamedDimsArraysTensorAlgebraExt QR (eltype=$(elt))" for elt in ( Float32, ComplexF32, Float64, ComplexF64 diff --git a/NDTensors/src/lib/NamedDimsArrays/test/test_basic.jl b/NDTensors/src/lib/NamedDimsArrays/test/test_basic.jl index 480e9a259a..1492255a74 100644 --- a/NDTensors/src/lib/NamedDimsArrays/test/test_basic.jl +++ b/NDTensors/src/lib/NamedDimsArrays/test/test_basic.jl @@ -10,10 +10,14 @@ using NDTensors.NamedDimsArrays: namedaxes, namedsize, unname +using NDTensors: NDTensors +using GPUArraysCore: @allowscalar +include(joinpath(pkgdir(NDTensors), "test", "NDTensorsTestUtils", "NDTensorsTestUtils.jl")) +using .NDTensorsTestUtils: devices_list @testset "NamedDimsArrays $(@__FILE__)" begin - @testset "Basic functionality" begin - a = randn(3, 4) + @testset "Basic functionality on $dev" for dev in devices_list(ARGS) + a = dev(randn(3, 4)) na = named(a, ("i", "j")) # TODO: Just call this `size`? i, j = namedsize(na) @@ -22,9 +26,11 @@ using NDTensors.NamedDimsArrays: @test !isnamed(a) @test isnamed(na) @test dimnames(na) == ("i", "j") - @test na[1, 1] == a[1, 1] - na[1, 1] = 11 - @test na[1, 1] == 11 + @allowscalar begin + @test na[1, 1] == a[1, 1] + na[1, 1] = 11 + @test na[1, 1] == 11 + end # TODO: Should `size` output `namedsize`? @test size(na) == (3, 4) @test namedsize(na) == (named(3, "i"), named(4, "j")) @@ -32,17 +38,19 @@ using NDTensors.NamedDimsArrays: # TODO: Should `axes` output `namedaxes`? @test axes(na) == (1:3, 1:4) @test namedaxes(na) == (named(1:3, "i"), named(1:4, "j")) - @test randn(named(3, "i"), named(4, "j")) isa NamedDimsArray - @test na["i" => 1, "j" => 2] == a[1, 2] - @test na["j" => 2, "i" => 1] == a[1, 2] - na["j" => 2, "i" => 1] = 12 - @test na[1, 2] == 12 - @test na[j => 1, i => 2] == a[2, 1] - @test na[aj => 1, ai => 2] == a[2, 1] - na[j => 1, i => 2] = 21 - @test na[2, 1] == 21 - na[aj => 1, ai => 2] = 2211 - @test na[2, 1] == 2211 + @test dev(randn(named(3, "i"), named(4, "j"))) isa NamedDimsArray + @allowscalar begin + @test na["i" => 1, "j" => 2] == a[1, 2] + @test na["j" => 2, "i" => 1] == a[1, 2] + na["j" => 2, "i" => 1] = 12 + @test na[1, 2] == 12 + @test na[j => 1, i => 2] == a[2, 1] + @test na[aj => 1, ai => 2] == a[2, 1] + na[j => 1, i => 2] = 21 + @test na[2, 1] == 21 + na[aj => 1, ai => 2] = 2211 + @test na[2, 1] == 2211 + end na′ = align(na, ("j", "i")) @test a == permutedims(unname(na′), (2, 1)) na′ = align(na, (j, i)) @@ -50,58 +58,64 @@ using NDTensors.NamedDimsArrays: na′ = align(na, (aj, ai)) @test a == permutedims(unname(na′), (2, 1)) end - @testset "Shorthand constructors (eltype=$elt)" for elt in ( - Float32, ComplexF32, Float64, ComplexF64 - ) + + @testset "Shorthand constructors (eltype=$elt)" for dev in devices_list(ARGS), + elt in (Float32, ComplexF32, Float64, ComplexF64) + i, j = named.((2, 2), ("i", "j")) value = rand(elt) - for na in (zeros(elt, i, j), zeros(elt, (i, j))) + for na in (dev(zeros(elt, i, j)), dev(zeros(elt, (i, j)))) @test eltype(na) === elt @test na isa NamedDimsArray @test dimnames(na) == ("i", "j") - @test iszero(na) + @allowscalar @test iszero(na) end - for na in (fill(value, i, j), fill(value, (i, j))) + for na in (dev(fill(value, i, j)), dev(fill(value, (i, j)))) @test eltype(na) === elt @test na isa NamedDimsArray @test dimnames(na) == ("i", "j") - @test all(isequal(value), na) + @allowscalar @test all(isequal(value), na) end - for na in (rand(elt, i, j), rand(elt, (i, j))) + for na in (dev(rand(elt, i, j)), dev(rand(elt, (i, j)))) @test eltype(na) === elt @test na isa NamedDimsArray @test dimnames(na) == ("i", "j") - @test !iszero(na) - @test all(x -> real(x) > 0, na) + @allowscalar begin + @test !iszero(na) + @test all(x -> real(x) > 0, na) + end end - for na in (randn(elt, i, j), randn(elt, (i, j))) + for na in (dev(randn(elt, i, j)), dev(randn(elt, (i, j)))) @test eltype(na) === elt @test na isa NamedDimsArray @test dimnames(na) == ("i", "j") - @test !iszero(na) + @allowscalar @test !iszero(na) end end - @testset "Shorthand constructors (eltype=unspecified)" begin + + @testset "Shorthand constructors (eltype=unspecified)" for dev in devices_list(ARGS) i, j = named.((2, 2), ("i", "j")) default_elt = Float64 - for na in (zeros(i, j), zeros((i, j))) + for na in (dev(zeros(i, j)), dev(zeros((i, j)))) @test eltype(na) === default_elt @test na isa NamedDimsArray @test dimnames(na) == ("i", "j") - @test iszero(na) + @allowscalar @test iszero(na) end - for na in (rand(i, j), rand((i, j))) + for na in (dev(rand(i, j)), dev(rand((i, j)))) @test eltype(na) === default_elt @test na isa NamedDimsArray @test dimnames(na) == ("i", "j") - @test !iszero(na) - @test all(x -> real(x) > 0, na) + @allowscalar begin + @test !iszero(na) + @test all(x -> real(x) > 0, na) + end end - for na in (randn(i, j), randn((i, j))) + for na in (dev(randn(i, j)), dev(randn((i, j)))) @test eltype(na) === default_elt @test na isa NamedDimsArray @test dimnames(na) == ("i", "j") - @test !iszero(na) + @allowscalar @test !iszero(na) end end end diff --git a/NDTensors/src/lib/TensorAlgebra/src/contract/allocate_output.jl b/NDTensors/src/lib/TensorAlgebra/src/contract/allocate_output.jl index f5935721fe..c219f8aa86 100644 --- a/NDTensors/src/lib/TensorAlgebra/src/contract/allocate_output.jl +++ b/NDTensors/src/lib/TensorAlgebra/src/contract/allocate_output.jl @@ -1,3 +1,4 @@ +using NDTensors.SetParameters: Position, set_parameter function allocate_output( ::typeof(contract), alg::Algorithm, @@ -13,7 +14,10 @@ function allocate_output( # TODO: Define `output_type(contract, alg, labels_dest, a1::AbstractArray, labels1, a2::AbstractArray, labels2, α, β)`. # TODO: Define `output_structure(contract, alg, labels_dest, a1::AbstractArray, labels1, a2::AbstractArray, labels2, α, β)`. # TODO: Define `allocate(type, structure)`. - return Array{promote_type(eltype(a1), eltype(a2))}(undef, length.(axes_dest)) + elt = promote_type(eltype(a1), eltype(a2)) + @assert typeof(a1) == typeof(a2) + array = set_parameter(typeof(a1), Position{1}(), elt) + return array(undef, length.(axes_dest)) end # TODO: Generalize to `output_structure`. diff --git a/NDTensors/src/lib/TensorAlgebra/src/contract/contract_matricize/contract.jl b/NDTensors/src/lib/TensorAlgebra/src/contract/contract_matricize/contract.jl index 1339287337..243564c832 100644 --- a/NDTensors/src/lib/TensorAlgebra/src/contract/contract_matricize/contract.jl +++ b/NDTensors/src/lib/TensorAlgebra/src/contract/contract_matricize/contract.jl @@ -1,3 +1,4 @@ +using NDTensors.Unwrap: expose function contract!( alg::Algorithm"matricize", a_dest::AbstractArray, @@ -12,7 +13,7 @@ function contract!( a_dest_matricized = matricize(a_dest, biperm_dest) a1_matricized = matricize(a1, biperm1) a2_matricized = matricize(a2, biperm2) - mul!(a_dest_matricized, a1_matricized, a2_matricized, α, β) + mul!(expose(a_dest_matricized), expose(a1_matricized), expose(a2_matricized), α, β) perm_dest = flatten(biperm_dest) # TODO: Create a function `unmatricize` or `unfusedims`. # unmatricize!(a_dest, a_dest_matricized, axes(a_dest), perm_dest) diff --git a/NDTensors/src/lib/UnallocatedArrays/README.md b/NDTensors/src/lib/UnallocatedArrays/README.md new file mode 100644 index 0000000000..f0e962a928 --- /dev/null +++ b/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. diff --git a/NDTensors/src/lib/UnallocatedArrays/src/UnallocatedArrays.jl b/NDTensors/src/lib/UnallocatedArrays/src/UnallocatedArrays.jl new file mode 100644 index 0000000000..ed2aa9d1b4 --- /dev/null +++ b/NDTensors/src/lib/UnallocatedArrays/src/UnallocatedArrays.jl @@ -0,0 +1,11 @@ +module UnallocatedArrays +include("abstractfill/abstractfill.jl") +include("abstractfill/set_types.jl") + +include("unallocatedfill.jl") +include("unallocatedzeros.jl") +include("abstractunallocatedarray.jl") +include("set_types.jl") + +export UnallocatedFill, UnallocatedZeros, alloctype, set_alloctype, allocate +end diff --git a/NDTensors/src/lib/UnallocatedArrays/src/abstractfill/abstractfill.jl b/NDTensors/src/lib/UnallocatedArrays/src/abstractfill/abstractfill.jl new file mode 100644 index 0000000000..7df2b799f3 --- /dev/null +++ b/NDTensors/src/lib/UnallocatedArrays/src/abstractfill/abstractfill.jl @@ -0,0 +1,19 @@ +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 + +allocate(A::AbstractFill) = alloctype(A)(parent(A)) + +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) diff --git a/NDTensors/src/lib/UnallocatedArrays/src/abstractfill/set_types.jl b/NDTensors/src/lib/UnallocatedArrays/src/abstractfill/set_types.jl new file mode 100644 index 0000000000..985e57885d --- /dev/null +++ b/NDTensors/src/lib/UnallocatedArrays/src/abstractfill/set_types.jl @@ -0,0 +1,75 @@ +using FillArrays: AbstractFill +using NDTensors.SetParameters: SetParameters, Position +using NDTensors.UnspecifiedTypes: UnspecifiedZero +## TODO make unit tests for all of these functions +## TODO remove P4 +# `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 +# right now I am just defining the necessary ones for my implementation still working on full implementation +# Set parameter 1 +SetParameters.set_parameter(T::Type{<:AbstractFill}, ::Position{1}, P1) = T{P1} + +# Set parameter 2 +function SetParameters.set_parameter( + T::Type{<:AbstractFill{P1}}, ::Position{2}, P2 +) where {P1} + return T{P2} +end + +# Set parameter 3 +function SetParameters.set_parameter( + T::Type{<:AbstractFill{P1,P2}}, ::Position{3}, P3 +) where {P1,P2} + return T{P3} +end + +## TODO define a specify_parameters function +## To quickly specify P1, P2, and P3 +## default parameters +function SetParameters.default_parameter(::Type{<:AbstractFill}, ::Position{1}) + return UnspecifiedTypes.UnspecifiedZero +end +SetParameters.default_parameter(::Type{<:AbstractFill}, ::Position{2}) = 0 +SetParameters.default_parameter(::Type{<:AbstractFill}, ::Position{3}) = Tuple{} + +SetParameters.nparameters(::Type{<:AbstractFill}) = Val(3) + +# Set parameter 1 +## Right now using AbstractArray +## TODO These are more difficult because T is technically defined so need some way to strip T of it {} types +# function set_parameter( +# T::Type{<:AbstractFill{<:Any,<:Any,P3}}, ::Position{1}, P1 +# ) where {P3} +# return T{P1,<:Any,P3} +# end +# function set_parameter( +# T::Type{<:AbstractFill{<:Any,P2,P3}}, ::Position{1}, P1 +# ) where {P2,P3} +# return T{P1,P2,P3} +# end +# function set_parameter( +# T::Type{<:AbstractFill{<:Any,<:Any,P3}}, ::Position{2}, P2 +# ) where {P3} +# return T{<:Any,P2,P3} +# end +# function set_parameter( +# T::Type{<:AbstractFill{P1,<:Any,P3}}, ::Position{2}, P2 +# ) where {P1,P3} +# return T +# end +# set_parameter(T::Type{<:AbstractFill}, ::Position{3}, P3) = T{<:Any,<:Any,P3} +# set_parameter(T::Type{<:AbstractFill{P1}}, ::Position{3}, P3) where {P1} = T{<:Any,P3} +# function set_parameter(T::Type{<:AbstractFill{<:Any,P2}}, ::Position{3}, P3) where {P2} +# return T{<:Any,P2,P3} +# end + +# Set parameter 2 +## using AbstractArray diff --git a/NDTensors/src/lib/UnallocatedArrays/src/abstractunallocatedarray.jl b/NDTensors/src/lib/UnallocatedArrays/src/abstractunallocatedarray.jl new file mode 100644 index 0000000000..8d981eb384 --- /dev/null +++ b/NDTensors/src/lib/UnallocatedArrays/src/abstractunallocatedarray.jl @@ -0,0 +1,52 @@ +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 set_alloctype( + complex(parent(A)), set_parameters(alloctype(A), Position{1}(), complex(eltype(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 Base.similar(a::UnallocatedArray) + return alloctype(a)(undef, size(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 + +## 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::) diff --git a/NDTensors/src/lib/UnallocatedArrays/src/set_types.jl b/NDTensors/src/lib/UnallocatedArrays/src/set_types.jl new file mode 100644 index 0000000000..6f4bfcbc34 --- /dev/null +++ b/NDTensors/src/lib/UnallocatedArrays/src/set_types.jl @@ -0,0 +1,38 @@ +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{<:UnallocatedFill{<:Any,<:Any,<:Any,P4}}, ::Position{4} +) where {P4} + return P4 +end +function SetParameters.get_parameter( + ::Type{<:UnallocatedZeros{<:Any,<:Any,<:Any,P4}}, ::Position{4} +) where {P4} + return P4 +end + +# ## Setting paramaters +function SetParameters.set_parameter( + T::Type{<:UnallocatedFill{P1,P2,P3,<:Any}}, ::Position{4}, P4::Type{<:AbstractArray} +) where {P1,P2,P3} + return T{P4} +end +function SetParameters.set_parameter( + T::Type{<:UnallocatedZeros{P1,P2,P3,<:Any}}, ::Position{4}, P4::Type{<:AbstractArray} +) where {P1,P2,P3} + return T{P4} +end + +# ## default parameters +function SetParameters.default_parameter(::Type{<:UnallocatedFill}, ::Position{4}) + return UnspecifiedArray{UnspecifiedNumber{UnspecifiedZero},0} +end +function SetParameters.default_parameter(::Type{<:UnallocatedZeros}, ::Position{4}) + return UnspecifiedArray{UnspecifiedNumber{UnspecifiedZero},0} +end + +SetParameters.nparameters(::Type{<:UnallocatedFill}) = Val(4) +SetParameters.nparameters(::Type{<:UnallocatedZeros}) = Val(4) diff --git a/NDTensors/src/lib/UnallocatedArrays/src/unallocatedfill.jl b/NDTensors/src/lib/UnallocatedArrays/src/unallocatedfill.jl new file mode 100644 index 0000000000..f98ade35d8 --- /dev/null +++ b/NDTensors/src/lib/UnallocatedArrays/src/unallocatedfill.jl @@ -0,0 +1,87 @@ +using FillArrays: FillArrays, AbstractFill, Fill, broadcasted_fill, kron_fill, mult_fill +using NDTensors.SetParameters: Position, set_parameters +import Base: + + +struct UnallocatedFill{ElT,N,Axes,Alloc} <: AbstractFill{ElT,N,Axes} + f::Fill{ElT,N,Axes} + alloc::Alloc + + function UnallocatedFill{ElT,N,Axes}( + f::Fill, alloc::Type{<:AbstractArray{ElT,N}} + ) where {ElT,N,Axes} + return new{ElT,N,Axes,Type{alloc}}(f, alloc) + end +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 + +function UnallocatedFill(f::Fill, alloc) + return UnallocatedFill{eltype(f)}(f, alloc) +end + +set_alloctype(f::Fill, alloc::Type{<:AbstractArray}) = 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 + +function +(A::UnallocatedFill, B::UnallocatedFill) + FillArrays.promote_shape(A, B) + b = getindex_value(B) + return A .+ b +end + +# #### +# ## TODO use `set_parameters` as constructor to these types +# function UnallocatedFill(f::Fill, alloc::Type{<:AbstractArray}) +# ## set_axes -> set_axes_type +# return set_alloctype( +# set_axes(set_ndims(set_eltype(UnallocatedFill, eltype(f)), ndims(f)), typeof(axes(f))), +# alloc, +# )( +# f +# ) +# end + +## Things to fix +## in different Change syntax of set_xxx_if_unspecified +# adapt diff --git a/NDTensors/src/lib/UnallocatedArrays/src/unallocatedzeros.jl b/NDTensors/src/lib/UnallocatedArrays/src/unallocatedzeros.jl new file mode 100644 index 0000000000..e6623ed71f --- /dev/null +++ b/NDTensors/src/lib/UnallocatedArrays/src/unallocatedzeros.jl @@ -0,0 +1,94 @@ +using FillArrays: + FillArrays, + AbstractZeros, + Fill, + Zeros, + broadcasted_fill, + broadcasted_zeros, + kron_fill, + kron_zeros, + mult_zeros +using NDTensors.SetParameters: Position, set_parameters +## TODO Should Alloc also be of ElT and N or should there be +## More freedom there? +struct UnallocatedZeros{ElT,N,Axes,Alloc} <: AbstractZeros{ElT,N,Axes} + z::Zeros{ElT,N,Axes} + alloc::Alloc + + function UnallocatedZeros{ElT,N,Axes}( + z::Zeros, alloc::Type{<:AbstractArray{ElT,N}} + ) where {ElT,N,Axes} + return new{ElT,N,Axes,Type{alloc}}(z, alloc) + end +end + +function UnallocatedZeros{ElT,N}(z::Zeros, alloc) where {ElT,N} + return UnallocatedZeros{ElT,N,typeof(axes(z))}(z, alloc) +end + +function UnallocatedZeros{ElT}(z::Zeros, alloc) where {ElT} + return UnallocatedZeros{ElT,ndims(z)}(z, alloc) +end + +function UnallocatedZeros(z::Zeros, alloc) + return UnallocatedZeros{eltype(z)}(z, alloc) +end + +set_alloctype(f::Zeros, alloc::Type{<:AbstractArray}) = UnallocatedZeros(f, alloc) + +Base.parent(Z::UnallocatedZeros) = Z.z + +Base.convert(::Type{<:UnallocatedZeros}, A::UnallocatedZeros) = A + +############################################# +# Arithmatic + +function FillArrays.mult_zeros(a::UnallocatedZeros, b, elt, ax) + return UnallocatedZeros(Zeros{elt}(ax), alloctype(a)) +end +FillArrays.mult_zeros(a, b::UnallocatedZeros, elt, ax) = mult_zeros(b, a, elt, ax) +function FillArrays.mult_zeros(a::UnallocatedZeros, b::UnallocatedZeros, elt, ax) + @assert(alloctype(a) == alloctype(b)) + return UnallocatedZeros(Zeros{elt}(ax), alloctype(a)) +end + +function FillArrays.broadcasted_zeros(f, a::UnallocatedZeros, elt, ax) + return UnallocatedZeros(Zeros{elt}(ax), alloctype(a)) +end +function FillArrays.broadcasted_zeros(f, a::UnallocatedZeros, b::UnallocatedZeros, elt, ax) + @assert(alloctype(a) == alloctype(b)) + return UnallocatedZeros(Zeros{elt}(ax), alloctype(a)) +end + +function FillArrays.broadcasted_zeros(f, a::UnallocatedZeros, b, elt, ax) + return UnallocatedZeros(Zeros{elt}(ax), alloctype(a)) +end +function FillArrays.broadcasted_zeros(f, a, b::UnallocatedZeros, elt, ax) + return broadcasted_zeros(f, b, a, elt, ax) +end + +function FillArrays.broadcasted_fill(f, a::UnallocatedZeros, val, ax) + return UnallocatedFill(Fill(val, ax), alloctype(a)) +end +function FillArrays.broadcasted_fill(f, a::UnallocatedZeros, b, val, ax) + return UnallocatedFill(Fill(val, ax), alloctype(a)) +end + +function FillArrays.broadcasted_fill(f, a, b::UnallocatedZeros, val, ax) + return broadcasted_fill(f, b, a, val, ax) +end + +function FillArrays.kron_zeros(a::UnallocatedZeros, b::UnallocatedZeros, elt, ax) + @assert alloctype(a) == alloctype(b) + return UnallocatedZeros(Zeros{elt}(ax), alloctype(a)) +end + +function FillArrays.kron_fill(a::UnallocatedZeros, b::UnallocatedFill, val, ax) + @assert alloctype(a) == alloctype(b) + elt = typeof(val) + return UnallocatedZeros(Zeros{elt}(ax), alloctype(a)) +end + +function FillArrays.kron_fill(a::UnallocatedFill, b::UnallocatedZeros, val, ax) + return kron_fill(b, a, val, ax) +end diff --git a/NDTensors/src/lib/UnallocatedArrays/test/runtests.jl b/NDTensors/src/lib/UnallocatedArrays/test/runtests.jl new file mode 100644 index 0000000000..f1236673d3 --- /dev/null +++ b/NDTensors/src/lib/UnallocatedArrays/test/runtests.jl @@ -0,0 +1,207 @@ +@eval module $(gensym()) +using FillArrays: FillArrays, AbstractFill, Fill, Zeros +using NDTensors: NDTensors +using NDTensors.UnallocatedArrays +using LinearAlgebra: norm +using Test: @test, @testset, @test_broken + +include(joinpath(pkgdir(NDTensors), "test", "NDTensorsTestUtils", "NDTensorsTestUtils.jl")) +using .NDTensorsTestUtils: devices_list + +@testset "Testing UnallocatedArrays" for dev in devices_list(ARGS), + elt in (Float64, Float32, ComplexF64, ComplexF32) + + @testset "Basic funcitonality" begin + z = Zeros{elt}((2, 3)) + Z = UnallocatedZeros(z, dev(Matrix{eltype(z)})) + + @test Z isa AbstractFill + @test size(Z) == (2, 3) + @test length(Z) == 6 + @test sum(Z) == 0 + @test norm(Z) == 0 + @test Z[2, 3] == 0 + @test allocate(Z) isa dev(Matrix{elt}) + Zp = set_alloctype(z, dev(Matrix{elt})) + @test Zp == Z + Zc = copy(Z) + @test Zc == Z + Zc = complex(Z) + @test eltype(Zc) == complex(eltype(z)) + @test Zc[1, 2] == 0.0 + 0.0im + + Zs = similar(Z) + @test Zs isa alloctype(Z) + + ######################################### + # UnallocatedFill + f = Fill{elt}(3.0, (2, 3, 4)) + F = UnallocatedFill(f, Array{elt,ndims(f)}) + @test F isa AbstractFill + @test size(F) == (2, 3, 4) + @test length(F) == 24 + @test sum(F) ≈ 3 * 24 + @test norm(F) ≈ sqrt(3^2 * 24) + @test F[2, 3, 1] == 3.0 + @test allocate(F) isa Array{elt,3} + Fp = allocate(F) + @test norm(Fp) ≈ norm(F) + Fs = similar(F) + @test Fs isa alloctype(F) + @test Fs[1, 1, 1] != 3.0 + + Fp = set_alloctype(f, dev(Array{elt,ndims(f)})) + @test allocate(Fp) isa dev(Array{elt,ndims(f)}) + @test Fp == F + Fc = copy(F) + @test Fc == F + Fc = allocate(complex(F)) + @test eltype(Fc) == complex(eltype(F)) + @test typeof(Fc) == alloctype(complex(F)) + Fc[2, 3, 4] = 4.0 + 3.0im + @test Fc[2, 3, 4] == 4.0 + 3.0im + ## TODO this isn't working + #Fc[2,3,4] = Base.setindex(F, 4.0+3.0im, 2,3,4) + end + + @testset "Multiplication" begin + z = Zeros{elt}((2, 3)) + Z = UnallocatedZeros(z, dev(Matrix{eltype(z)})) + + R = Z * Z' + @test R isa UnallocatedZeros + @test alloctype(R) == alloctype(Z) + @test size(R) == (2, 2) + M = rand(elt, (3, 4)) + R = Z * M + @test R isa UnallocatedZeros + @test alloctype(R) == alloctype(Z) + @test size(R) == (2, 4) + R = M' * Z' + @test R isa UnallocatedZeros + @test alloctype(R) == alloctype(Z) + @test size(R) == (4, 2) + R = transpose(M) * transpose(Z) + @test R isa UnallocatedZeros + @test alloctype(R) == alloctype(Z) + @test size(R) == (4, 2) + + ################################### + ## UnallocatedFill + f = Fill{elt}(3.0, (2, 12)) + F = UnallocatedFill(f, dev(Matrix{elt})) + p = Fill{elt}(4.0, (12, 5)) + P = UnallocatedFill(p, dev(Array{elt,ndims(p)})) + R = F * P + @test F isa UnallocatedFill + @test R[1, 1] == 144 + @test alloctype(R) == alloctype(F) + @test size(R) == (2, 5) + + R = F * F' + @test R isa UnallocatedFill + @test R[1, 2] == elt(108) + @test alloctype(R) == alloctype(F) + @test size(R) == (2, 2) + + R = transpose(F) * F + @test R isa UnallocatedFill + @test R[12, 3] == elt(18) + @test alloctype(R) == alloctype(F) + @test size(R) == (12, 12) + + R = transpose(Z) * F + @test R isa UnallocatedZeros + @test alloctype(R) == alloctype(Z) + @test size(R) == (3, 12) + end + + @testset "Broadcast" begin + z = Zeros{elt}((2, 3)) + Z = UnallocatedZeros(z, dev(Matrix{elt})) + R = elt(2.0) .* Z + @test R isa UnallocatedZeros + @test alloctype(R) == alloctype(Z) + R = Z .* elt(2.0) + @test R isa UnallocatedZeros + @test alloctype(R) == alloctype(Z) + + R = Z .* Z + @test R isa UnallocatedZeros + @test alloctype(R) == alloctype(Z) + + Z = UnallocatedZeros(Zeros{elt}((2, 3)), dev(Matrix{elt})) + R = Z + Z + @test R isa UnallocatedZeros + @test alloctype(R) == alloctype(Z) + + R = Z .+ elt(2.0) + @test R isa UnallocatedFill + @test alloctype(R) == alloctype(Z) + ######################## + # UnallocatedFill + f = Fill(elt(3.0), (2, 3, 4)) + F = UnallocatedFill(f, Array{elt,ndims(f)}) + F2 = F .* 2 + @test F2 isa UnallocatedFill + @test F2[1, 1, 1] == elt(6.0) + @test alloctype(F2) == alloctype(F) + + #F2 .+= elt(2.0) ## This is broken + F2 = F2 .+ elt(2.0) + @test F2 isa UnallocatedFill + @test F2[1, 1, 1] == elt(8.0) + @test alloctype(F2) == alloctype(F) + + F = UnallocatedFill(Fill(elt(2.0), (2, 3)), dev(Matrix{elt})) + R = Z + F + @test R isa UnallocatedFill + @test alloctype(R) == alloctype(Z) + + R = F + Z + @test R isa UnallocatedFill + @test alloctype(R) == alloctype(Z) + + F = UnallocatedFill(Fill(elt(3.0), (2, 12)), dev(Matrix{elt})) + R = F .* F + @test R isa UnallocatedFill + @test R[2, 9] == elt(9) + @test alloctype(R) == alloctype(F) + @test size(R) == (2, 12) + + P = UnallocatedFill(Fill(elt(4.0), (2, 3)), dev(Matrix{elt})) + R = Z .* P + @test R isa UnallocatedZeros + @test alloctype(R) == alloctype(P) + @test size(R) == (2, 3) + + F = UnallocatedFill(Fill(elt(2), (2, 3)), dev(Matrix{elt})) + R = F + F + @test R isa UnallocatedFill + end + + ## TODO make other kron tests + @testset "Kron" begin + A = UnallocatedZeros(Zeros{elt}(2), dev(Vector{elt})) + B = UnallocatedZeros(Zeros{elt}(2), dev(Vector{elt})) + C = kron(A, B) + @test C isa UnallocatedZeros + @test alloctype(C) == alloctype(B) + + B = UnallocatedFill(Fill(elt(2.0), (2)), dev(Vector{elt})) + C = kron(A, B) + @test C isa UnallocatedZeros + @test alloctype(C) == alloctype(B) + + C = kron(B, A) + @test C isa UnallocatedZeros + @test alloctype(C) == alloctype(B) + + A = UnallocatedFill(Fill(elt(3.0), (2)), dev(Vector{elt})) + C = kron(A, B) + @test C isa UnallocatedFill + @test alloctype(C) == alloctype(B) + @test C[1] == elt(6) + end +end +end diff --git a/NDTensors/src/lib/UnspecifiedTypes/README.md b/NDTensors/src/lib/UnspecifiedTypes/README.md new file mode 100644 index 0000000000..15ad85367d --- /dev/null +++ b/NDTensors/src/lib/UnspecifiedTypes/README.md @@ -0,0 +1,3 @@ +# UnspecifiedTypes + +A module defining a set of basic types which are place holders for allocated bit-wise representable types. diff --git a/NDTensors/src/lib/UnspecifiedTypes/src/UnspecifiedTypes.jl b/NDTensors/src/lib/UnspecifiedTypes/src/UnspecifiedTypes.jl new file mode 100644 index 0000000000..f94a63a14c --- /dev/null +++ b/NDTensors/src/lib/UnspecifiedTypes/src/UnspecifiedTypes.jl @@ -0,0 +1,11 @@ +module UnspecifiedTypes + +using LinearAlgebra + +include("unspecifiednumber.jl") +include("unspecifiedzero.jl") + +include("unspecifiedarray.jl") + +export UnspecifiedArray, UnspecifiedNumber, UnspecifiedZero +end diff --git a/NDTensors/src/lib/UnspecifiedTypes/src/unspecifiedarray.jl b/NDTensors/src/lib/UnspecifiedTypes/src/unspecifiedarray.jl new file mode 100644 index 0000000000..e47e9832e5 --- /dev/null +++ b/NDTensors/src/lib/UnspecifiedTypes/src/unspecifiedarray.jl @@ -0,0 +1 @@ +struct UnspecifiedArray{ElT,N} <: AbstractArray{ElT,N} end diff --git a/NDTensors/src/lib/UnspecifiedTypes/src/unspecifiednumber.jl b/NDTensors/src/lib/UnspecifiedTypes/src/unspecifiednumber.jl new file mode 100644 index 0000000000..c44587b376 --- /dev/null +++ b/NDTensors/src/lib/UnspecifiedTypes/src/unspecifiednumber.jl @@ -0,0 +1,5 @@ +abstract type AbstractUnspecifiedNumber <: Number end + +struct UnspecifiedNumber{T} <: AbstractUnspecifiedNumber + value::T +end diff --git a/NDTensors/src/lib/UnspecifiedTypes/src/unspecifiedzero.jl b/NDTensors/src/lib/UnspecifiedTypes/src/unspecifiedzero.jl new file mode 100644 index 0000000000..4ca1858218 --- /dev/null +++ b/NDTensors/src/lib/UnspecifiedTypes/src/unspecifiedzero.jl @@ -0,0 +1,39 @@ +struct UnspecifiedZero <: AbstractUnspecifiedNumber end + +# Base.Complex{UnspecifiedZero}() = complex(UnspecifiedZero()) +# function Base.Complex{UnspecifiedZero}(z::Real) +# return (iszero(z) ? complex(UnspecifiedZero()) : throw(ErrorException)) +# end + +zero(::Type{UnspecifiedZero}) = UnspecifiedZero() +zero(n::UnspecifiedZero) = zero(typeof(n)) + +# This helps handle a lot of basic algebra, like: +# UnspecifiedZero() + 2.3 == 2.3 +convert(::Type{T}, x::UnspecifiedZero) where {T<:Number} = T(zero(T)) + +#convert(::Type{Complex{UnspecifiedZero}}, x::UnspecifiedZero) = complex(x) + +# TODO: Should this be implemented? +#Complex(x::Real, ::UnspecifiedZero) = x + +# This is to help define `float(::UnspecifiedZero) = 0.0`. +# This helps with defining `norm` of `UnallocatedZeros{UnspecifiedZero}`. +AbstractFloat(::UnspecifiedZero) = zero(AbstractFloat) + +# Basic arithmetic +(::UnspecifiedZero + ::UnspecifiedZero) = UnspecifiedZero() +(::UnspecifiedZero - ::UnspecifiedZero) = UnspecifiedZero() +(::Number * ::UnspecifiedZero) = UnspecifiedZero() +(::UnspecifiedZero * ::Number) = UnspecifiedZero() +(::UnspecifiedZero * ::UnspecifiedZero) = UnspecifiedZero() +(::UnspecifiedZero / ::Number) = UnspecifiedZero() +(::Number / ::UnspecifiedZero) = throw(DivideError()) +(::UnspecifiedZero / ::UnspecifiedZero) = throw(DivideError()) +-(::UnspecifiedZero) = UnspecifiedZero() + +Base.promote_type(z::Type{<:UnspecifiedZero}, ElT::Type) = Base.promote_type(ElT, z) + +Base.promote_type(ElT::Type, ::Type{<:UnspecifiedZero}) = ElT +Base.promote_type(::Type{<:UnspecifiedZero}, ::Type{<:UnspecifiedZero}) = UnspecifiedZero +Base.promote_type(ElT::Type, ::Type{<:Complex{<:UnspecifiedZero}}) = Complex{real(ElT)} diff --git a/NDTensors/src/lib/UnspecifiedTypes/test/runtests.jl b/NDTensors/src/lib/UnspecifiedTypes/test/runtests.jl new file mode 100644 index 0000000000..700ad9fac2 --- /dev/null +++ b/NDTensors/src/lib/UnspecifiedTypes/test/runtests.jl @@ -0,0 +1,9 @@ +## TODO This needs work here +@eval module $(gensym()) +using NDTensors.UnspecifiedTypes +using Test: @test, @testset + +@testset "Testing UnspecifiedTypes" begin + UA = UnspecifiedArray{} +end +end diff --git a/NDTensors/test/Project.toml b/NDTensors/test/Project.toml index f47f1f7f21..f6ec19f498 100644 --- a/NDTensors/test/Project.toml +++ b/NDTensors/test/Project.toml @@ -7,6 +7,7 @@ Compat = "34da2185-b29b-5c13-b0c7-acf172513d20" Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" ITensors = "9136182c-28ba-11e9-034c-db9fb085ebd5" +FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" NDTensors = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf" Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4" diff --git a/NDTensors/test/lib/Project.toml b/NDTensors/test/lib/Project.toml index b86a01c764..41bb43fa99 100644 --- a/NDTensors/test/lib/Project.toml +++ b/NDTensors/test/lib/Project.toml @@ -4,6 +4,7 @@ BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" Compat = "34da2185-b29b-5c13-b0c7-acf172513d20" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" +FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" NDTensors = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2" diff --git a/NDTensors/test/lib/runtests.jl b/NDTensors/test/lib/runtests.jl index d01ccfe126..14c42ee3fc 100644 --- a/NDTensors/test/lib/runtests.jl +++ b/NDTensors/test/lib/runtests.jl @@ -15,6 +15,8 @@ using Test: @testset "SparseArrayDOKs", "TagSets", "TensorAlgebra", + "UnallocatedArrays", + "UnspecifiedTypes", "Unwrap", ] using NDTensors: NDTensors