changeset 325:41c3c25e4e3b

LazyTensors: Simplify the LazyElementwiseOperation type by restricting it and introducing a LazyConstantArray to handle scalars.
author Jonatan Werpers <jonatan@werpers.com>
date Thu, 24 Sep 2020 22:31:04 +0200
parents 047dee8efaef
children c8bbb4a83518
files LazyTensors/src/lazy_array.jl LazyTensors/test/runtests.jl
diffstat 2 files changed, 43 insertions(+), 104 deletions(-) [+]
line wrap: on
line diff
--- a/LazyTensors/src/lazy_array.jl	Thu Sep 24 21:28:28 2020 +0200
+++ b/LazyTensors/src/lazy_array.jl	Thu Sep 24 22:31:04 2020 +0200
@@ -8,124 +8,54 @@
 abstract type LazyArray{T,D} <: AbstractArray{T,D} end
 export LazyArray
 
+struct LazyConstantArray{T,D} <: LazyArray{T,D}
+	val::T
+	size::NTuple{D,Int}
+end
+
+Base.size(lca::LazyConstantArray) = lca.size
+Base.getindex(lca::LazyConstantArray{T,D}, I::Vararg{Int,D}) where {T,D} = lca.val
+
 """
-    LazyElementwiseOperation{T,D,Op,T1,T2} <: LazyArray{T,D}
+    LazyElementwiseOperation{T,D,Op} <: LazyArray{T,D}
 Struct allowing for lazy evaluation of elementwise operations on AbstractArrays.
 
-A LazyElementwiseOperation contains two datatypes T1, and T2, together with an operation,
-where at least one of T1 and T2 is an AbstractArray, and one may be a Real.
+A LazyElementwiseOperation contains two arrays together with an operation.
 The operations are carried out when the LazyElementwiseOperation is indexed.
 """
-struct LazyElementwiseOperation{T,D,Op,T1,T2} <: LazyArray{T,D}
-    a::T1
-    b::T2
+struct LazyElementwiseOperation{T,D,Op} <: LazyArray{T,D}
+    a::AbstractArray{T,D}
+    b::AbstractArray{T,D}
 
-    @inline function LazyElementwiseOperation{T,D,Op}(a::T1,b::T2) where {T,D,Op,T1<:AbstractArray{T,D},T2<:AbstractArray{T,D}}
+    function LazyElementwiseOperation{T,D,Op}(a::AbstractArray{T,D},b::AbstractArray{T,D}) where {T,D,Op}
         @boundscheck if size(a) != size(b)
             throw(DimensionMismatch("dimensions must match"))
         end
-        return new{T,D,Op,T1,T2}(a,b)
+        return new{T,D,Op}(a,b)
     end
 
-    @inline function LazyElementwiseOperation{T,D,Op}(a::T1,b::T2) where {T,D,Op,T1<:AbstractArray{T,D},T2<:Real}
-        return new{T,D,Op,T1,T2}(a,b)
-    end
-
-    @inline function LazyElementwiseOperation{T,D,Op}(a::T1,b::T2) where {T,D,Op,T1<:Real,T2<:AbstractArray{T,D}}
-        return new{T,D,Op,T1,T2}(a,b)
-    end
+	LazyElementwiseOperation{T,D,Op}(a::AbstractArray{T,D},b::T) where {T,D,Op} = new{T,D,Op}(a, LazyConstantArray(b, size(a)))
+	LazyElementwiseOperation{T,D,Op}(a::T,b::AbstractArray{T,D}) where {T,D,Op} = new{T,D,Op}(LazyConstantArray(a,  size(b)), b)
 end
 # TODO: Move Op to be the first parameter? Compare to Binary operations
 
 Base.size(v::LazyElementwiseOperation) = size(v.a)
 
+evaluate(leo::LazyElementwiseOperation{T,D,:+}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] + leo.b[I...]
+evaluate(leo::LazyElementwiseOperation{T,D,:-}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] - leo.b[I...]
+evaluate(leo::LazyElementwiseOperation{T,D,:*}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] * leo.b[I...]
+evaluate(leo::LazyElementwiseOperation{T,D,:/}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] / leo.b[I...]
+
 # TODO: Make sure boundschecking is done properly and that the lenght of the vectors are equal
 # NOTE: Boundschecking in getindex functions now assumes that the size of the
 # vectors in the LazyElementwiseOperation are the same size. If we remove the
 # size assertion in the constructor we might have to handle
 # boundschecking differently.
-Base.@propagate_inbounds @inline function Base.getindex(leo::LazyElementwiseOperation{T,D,:+,T1,T2}, I::Vararg{Int,D}) where {T,D,T1<:AbstractArray{T,D},T2<:AbstractArray{T,D}}
-    @boundscheck if !checkbounds(Bool,leo.a,I...)
-        throw(BoundsError([leo],I...))
-    end
-    return leo.a[I...] + leo.b[I...]
-end
-
-Base.@propagate_inbounds @inline function Base.getindex(leo::LazyElementwiseOperation{T,D,:-,T1,T2}, I::Vararg{Int,D}) where {T,D,T1<:AbstractArray{T,D},T2<:AbstractArray{T,D}}
-    @boundscheck if !checkbounds(Bool,leo.a,I...)
-        throw(BoundsError([leo],I...))
-    end
-    return leo.a[I...] - leo.b[I...]
-end
-
-Base.@propagate_inbounds @inline function Base.getindex(leo::LazyElementwiseOperation{T,D,:*,T1,T2}, I::Vararg{Int,D}) where {T,D,T1<:AbstractArray{T,D},T2<:AbstractArray{T,D}}
-    @boundscheck if !checkbounds(Bool,leo.a,I...)
-        throw(BoundsError([leo],I...))
-    end
-    return leo.a[I...] * leo.b[I...]
-end
-
-Base.@propagate_inbounds @inline function Base.getindex(leo::LazyElementwiseOperation{T,D,:/,T1,T2}, I::Vararg{Int,D}) where {T,D,T1<:AbstractArray{T,D},T2<:AbstractArray{T,D}}
-    @boundscheck if !checkbounds(Bool,leo.a,I...)
-        throw(BoundsError([leo],I...))
-    end
-    return leo.a[I...] / leo.b[I...]
-end
-
-Base.@propagate_inbounds @inline function Base.getindex(leo::LazyElementwiseOperation{T,D,:+,T1,T2}, I::Vararg{Int,D}) where {T,D,T1<:AbstractArray{T,D},T2<:Real}
-    @boundscheck if !checkbounds(Bool,leo.a,I...)
-        throw(BoundsError([leo],I...))
-    end
-    return leo.a[I...] + leo.b
-end
-
-Base.@propagate_inbounds @inline function Base.getindex(leo::LazyElementwiseOperation{T,D,:-,T1,T2}, I::Vararg{Int,D}) where {T,D,T1<:AbstractArray{T,D},T2<:Real}
-    @boundscheck if !checkbounds(Bool,leo.a,I...)
-        throw(BoundsError([leo],I...))
+Base.@propagate_inbounds @inline function Base.getindex(leo::LazyElementwiseOperation{T,D}, I::Vararg{Int,D}) where {T,D}
+    @boundscheck if !checkbounds(Bool, leo.a, I...)
+        throw(BoundsError([leo], I...))
     end
-    return leo.a[I...] - leo.b
-end
-
-Base.@propagate_inbounds @inline function Base.getindex(leo::LazyElementwiseOperation{T,D,:*,T1,T2}, I::Vararg{Int,D}) where {T,D,T1<:AbstractArray{T,D},T2<:Real}
-    @boundscheck if !checkbounds(Bool,leo.a,I...)
-        throw(BoundsError([leo],I...))
-    end
-    return leo.a[I...] * leo.b
-end
-
-Base.@propagate_inbounds @inline function Base.getindex(leo::LazyElementwiseOperation{T,D,:/,T1,T2}, I::Vararg{Int,D}) where {T,D,T1<:AbstractArray{T,D},T2<:Real}
-    @boundscheck if !checkbounds(Bool,leo.a,I...)
-        throw(BoundsError([leo],I...))
-    end
-    return leo.a[I...] / leo.b
-end
-
-Base.@propagate_inbounds @inline function Base.getindex(leo::LazyElementwiseOperation{T,D,:+,T1,T2}, I::Vararg{Int,D}) where {T,D,T1<:Real,T2<:AbstractArray{T,D}}
-    @boundscheck if !checkbounds(Bool,leo.b,I...)
-        throw(BoundsError([leo],I...))
-    end
-    return leo.a + leo.b[I...]
-end
-
-Base.@propagate_inbounds @inline function Base.getindex(leo::LazyElementwiseOperation{T,D,:-,T1,T2}, I::Vararg{Int,D}) where {T,D,T1<:Real,T2<:AbstractArray{T,D}}
-    @boundscheck if !checkbounds(Bool,leo.b,I...)
-        throw(BoundsError([leo],I...))
-    end
-    return leo.a - leo.b[I...]
-end
-
-Base.@propagate_inbounds @inline function Base.getindex(leo::LazyElementwiseOperation{T,D,:*,T1,T2}, I::Vararg{Int,D}) where {T,D,T1<:Real,T2<:AbstractArray{T,D}}
-    @boundscheck if !checkbounds(Bool,leo.b,I...)
-        throw(BoundsError([leo],I...))
-    end
-    return leo.a * leo.b[I...]
-end
-
-Base.@propagate_inbounds @inline function Base.getindex(leo::LazyElementwiseOperation{T,D,:/,T1,T2}, I::Vararg{Int,D}) where {T,D,T1<:Real,T2<:AbstractArray{T,D}}
-    @boundscheck if !checkbounds(Bool,leo.b,I...)
-        throw(BoundsError([leo],I...))
-    end
-    return leo.a / leo.b[I...]
+    return evaluate(leo, I...)
 end
 
 # Define lazy operations for AbstractArrays. Operations constructs a LazyElementwiseOperation which
@@ -135,15 +65,15 @@
 Base.@propagate_inbounds *̃(a::AbstractArray{T,D}, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:*}(a,b)
 Base.@propagate_inbounds /̃(a::AbstractArray{T,D}, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:/}(a,b)
 
-Base.@propagate_inbounds +̃(a::AbstractArray{T,D}, b::Real) where {T,D} = LazyElementwiseOperation{T,D,:+}(a,b)
-Base.@propagate_inbounds -̃(a::AbstractArray{T,D}, b::Real) where {T,D} = LazyElementwiseOperation{T,D,:-}(a,b)
-Base.@propagate_inbounds *̃(a::AbstractArray{T,D}, b::Real) where {T,D} = LazyElementwiseOperation{T,D,:*}(a,b)
-Base.@propagate_inbounds /̃(a::AbstractArray{T,D}, b::Real) where {T,D} = LazyElementwiseOperation{T,D,:/}(a,b)
+Base.@propagate_inbounds +̃(a::AbstractArray{T,D}, b::T) where {T,D} = LazyElementwiseOperation{T,D,:+}(a,b)
+Base.@propagate_inbounds -̃(a::AbstractArray{T,D}, b::T) where {T,D} = LazyElementwiseOperation{T,D,:-}(a,b)
+Base.@propagate_inbounds *̃(a::AbstractArray{T,D}, b::T) where {T,D} = LazyElementwiseOperation{T,D,:*}(a,b)
+Base.@propagate_inbounds /̃(a::AbstractArray{T,D}, b::T) where {T,D} = LazyElementwiseOperation{T,D,:/}(a,b)
 
-Base.@propagate_inbounds +̃(a::Real, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:+}(a,b)
-Base.@propagate_inbounds -̃(a::Real, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:-}(a,b)
-Base.@propagate_inbounds *̃(a::Real, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:*}(a,b)
-Base.@propagate_inbounds /̃(a::Real, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:/}(a,b)
+Base.@propagate_inbounds +̃(a::T, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:+}(a,b)
+Base.@propagate_inbounds -̃(a::T, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:-}(a,b)
+Base.@propagate_inbounds *̃(a::T, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:*}(a,b)
+Base.@propagate_inbounds /̃(a::T, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:/}(a,b)
 
 
 
--- a/LazyTensors/test/runtests.jl	Thu Sep 24 21:28:28 2020 +0200
+++ b/LazyTensors/test/runtests.jl	Thu Sep 24 22:31:04 2020 +0200
@@ -115,6 +115,15 @@
 end
 
 @testset "LazyArray" begin
+	@testset "LazyConstantArray" begin
+	    @test LazyTensors.LazyConstantArray(3,(3,2)) isa LazyArray{Int,2}
+
+	    lca = LazyTensors.LazyConstantArray(3.0,(3,2))
+	    @test eltype(lca) == Float64
+	    @test ndims(lca) == 2
+	    @test size(lca) == (3,2)
+	    @test lca[2] == 3.0
+	end
     struct DummyArray{T,D, T1<:AbstractArray{T,D}} <: LazyArray{T,D}
         data::T1
     end