comparison LazyTensors/src/lazy_array.jl @ 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 634453a4e1d8
children
comparison
equal deleted inserted replaced
324:047dee8efaef 325:41c3c25e4e3b
6 A subtype of `LazyArray` will use lazy version of `+`, `-`, `*`, `/`. 6 A subtype of `LazyArray` will use lazy version of `+`, `-`, `*`, `/`.
7 """ 7 """
8 abstract type LazyArray{T,D} <: AbstractArray{T,D} end 8 abstract type LazyArray{T,D} <: AbstractArray{T,D} end
9 export LazyArray 9 export LazyArray
10 10
11 struct LazyConstantArray{T,D} <: LazyArray{T,D}
12 val::T
13 size::NTuple{D,Int}
14 end
15
16 Base.size(lca::LazyConstantArray) = lca.size
17 Base.getindex(lca::LazyConstantArray{T,D}, I::Vararg{Int,D}) where {T,D} = lca.val
18
11 """ 19 """
12 LazyElementwiseOperation{T,D,Op,T1,T2} <: LazyArray{T,D} 20 LazyElementwiseOperation{T,D,Op} <: LazyArray{T,D}
13 Struct allowing for lazy evaluation of elementwise operations on AbstractArrays. 21 Struct allowing for lazy evaluation of elementwise operations on AbstractArrays.
14 22
15 A LazyElementwiseOperation contains two datatypes T1, and T2, together with an operation, 23 A LazyElementwiseOperation contains two arrays together with an operation.
16 where at least one of T1 and T2 is an AbstractArray, and one may be a Real.
17 The operations are carried out when the LazyElementwiseOperation is indexed. 24 The operations are carried out when the LazyElementwiseOperation is indexed.
18 """ 25 """
19 struct LazyElementwiseOperation{T,D,Op,T1,T2} <: LazyArray{T,D} 26 struct LazyElementwiseOperation{T,D,Op} <: LazyArray{T,D}
20 a::T1 27 a::AbstractArray{T,D}
21 b::T2 28 b::AbstractArray{T,D}
22 29
23 @inline function LazyElementwiseOperation{T,D,Op}(a::T1,b::T2) where {T,D,Op,T1<:AbstractArray{T,D},T2<:AbstractArray{T,D}} 30 function LazyElementwiseOperation{T,D,Op}(a::AbstractArray{T,D},b::AbstractArray{T,D}) where {T,D,Op}
24 @boundscheck if size(a) != size(b) 31 @boundscheck if size(a) != size(b)
25 throw(DimensionMismatch("dimensions must match")) 32 throw(DimensionMismatch("dimensions must match"))
26 end 33 end
27 return new{T,D,Op,T1,T2}(a,b) 34 return new{T,D,Op}(a,b)
28 end 35 end
29 36
30 @inline function LazyElementwiseOperation{T,D,Op}(a::T1,b::T2) where {T,D,Op,T1<:AbstractArray{T,D},T2<:Real} 37 LazyElementwiseOperation{T,D,Op}(a::AbstractArray{T,D},b::T) where {T,D,Op} = new{T,D,Op}(a, LazyConstantArray(b, size(a)))
31 return new{T,D,Op,T1,T2}(a,b) 38 LazyElementwiseOperation{T,D,Op}(a::T,b::AbstractArray{T,D}) where {T,D,Op} = new{T,D,Op}(LazyConstantArray(a, size(b)), b)
32 end
33
34 @inline function LazyElementwiseOperation{T,D,Op}(a::T1,b::T2) where {T,D,Op,T1<:Real,T2<:AbstractArray{T,D}}
35 return new{T,D,Op,T1,T2}(a,b)
36 end
37 end 39 end
38 # TODO: Move Op to be the first parameter? Compare to Binary operations 40 # TODO: Move Op to be the first parameter? Compare to Binary operations
39 41
40 Base.size(v::LazyElementwiseOperation) = size(v.a) 42 Base.size(v::LazyElementwiseOperation) = size(v.a)
43
44 evaluate(leo::LazyElementwiseOperation{T,D,:+}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] + leo.b[I...]
45 evaluate(leo::LazyElementwiseOperation{T,D,:-}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] - leo.b[I...]
46 evaluate(leo::LazyElementwiseOperation{T,D,:*}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] * leo.b[I...]
47 evaluate(leo::LazyElementwiseOperation{T,D,:/}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] / leo.b[I...]
41 48
42 # TODO: Make sure boundschecking is done properly and that the lenght of the vectors are equal 49 # TODO: Make sure boundschecking is done properly and that the lenght of the vectors are equal
43 # NOTE: Boundschecking in getindex functions now assumes that the size of the 50 # NOTE: Boundschecking in getindex functions now assumes that the size of the
44 # vectors in the LazyElementwiseOperation are the same size. If we remove the 51 # vectors in the LazyElementwiseOperation are the same size. If we remove the
45 # size assertion in the constructor we might have to handle 52 # size assertion in the constructor we might have to handle
46 # boundschecking differently. 53 # boundschecking differently.
47 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}} 54 Base.@propagate_inbounds @inline function Base.getindex(leo::LazyElementwiseOperation{T,D}, I::Vararg{Int,D}) where {T,D}
48 @boundscheck if !checkbounds(Bool,leo.a,I...) 55 @boundscheck if !checkbounds(Bool, leo.a, I...)
49 throw(BoundsError([leo],I...)) 56 throw(BoundsError([leo], I...))
50 end 57 end
51 return leo.a[I...] + leo.b[I...] 58 return evaluate(leo, I...)
52 end
53
54 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}}
55 @boundscheck if !checkbounds(Bool,leo.a,I...)
56 throw(BoundsError([leo],I...))
57 end
58 return leo.a[I...] - leo.b[I...]
59 end
60
61 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}}
62 @boundscheck if !checkbounds(Bool,leo.a,I...)
63 throw(BoundsError([leo],I...))
64 end
65 return leo.a[I...] * leo.b[I...]
66 end
67
68 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}}
69 @boundscheck if !checkbounds(Bool,leo.a,I...)
70 throw(BoundsError([leo],I...))
71 end
72 return leo.a[I...] / leo.b[I...]
73 end
74
75 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}
76 @boundscheck if !checkbounds(Bool,leo.a,I...)
77 throw(BoundsError([leo],I...))
78 end
79 return leo.a[I...] + leo.b
80 end
81
82 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}
83 @boundscheck if !checkbounds(Bool,leo.a,I...)
84 throw(BoundsError([leo],I...))
85 end
86 return leo.a[I...] - leo.b
87 end
88
89 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}
90 @boundscheck if !checkbounds(Bool,leo.a,I...)
91 throw(BoundsError([leo],I...))
92 end
93 return leo.a[I...] * leo.b
94 end
95
96 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}
97 @boundscheck if !checkbounds(Bool,leo.a,I...)
98 throw(BoundsError([leo],I...))
99 end
100 return leo.a[I...] / leo.b
101 end
102
103 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}}
104 @boundscheck if !checkbounds(Bool,leo.b,I...)
105 throw(BoundsError([leo],I...))
106 end
107 return leo.a + leo.b[I...]
108 end
109
110 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}}
111 @boundscheck if !checkbounds(Bool,leo.b,I...)
112 throw(BoundsError([leo],I...))
113 end
114 return leo.a - leo.b[I...]
115 end
116
117 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}}
118 @boundscheck if !checkbounds(Bool,leo.b,I...)
119 throw(BoundsError([leo],I...))
120 end
121 return leo.a * leo.b[I...]
122 end
123
124 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}}
125 @boundscheck if !checkbounds(Bool,leo.b,I...)
126 throw(BoundsError([leo],I...))
127 end
128 return leo.a / leo.b[I...]
129 end 59 end
130 60
131 # Define lazy operations for AbstractArrays. Operations constructs a LazyElementwiseOperation which 61 # Define lazy operations for AbstractArrays. Operations constructs a LazyElementwiseOperation which
132 # can later be indexed into. Lazy operations are denoted by the usual operator followed by a tilde 62 # can later be indexed into. Lazy operations are denoted by the usual operator followed by a tilde
133 Base.@propagate_inbounds +̃(a::AbstractArray{T,D}, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:+}(a,b) 63 Base.@propagate_inbounds +̃(a::AbstractArray{T,D}, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:+}(a,b)
134 Base.@propagate_inbounds -̃(a::AbstractArray{T,D}, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:-}(a,b) 64 Base.@propagate_inbounds -̃(a::AbstractArray{T,D}, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:-}(a,b)
135 Base.@propagate_inbounds *̃(a::AbstractArray{T,D}, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:*}(a,b) 65 Base.@propagate_inbounds *̃(a::AbstractArray{T,D}, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:*}(a,b)
136 Base.@propagate_inbounds /̃(a::AbstractArray{T,D}, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:/}(a,b) 66 Base.@propagate_inbounds /̃(a::AbstractArray{T,D}, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:/}(a,b)
137 67
138 Base.@propagate_inbounds +̃(a::AbstractArray{T,D}, b::Real) where {T,D} = LazyElementwiseOperation{T,D,:+}(a,b) 68 Base.@propagate_inbounds +̃(a::AbstractArray{T,D}, b::T) where {T,D} = LazyElementwiseOperation{T,D,:+}(a,b)
139 Base.@propagate_inbounds -̃(a::AbstractArray{T,D}, b::Real) where {T,D} = LazyElementwiseOperation{T,D,:-}(a,b) 69 Base.@propagate_inbounds -̃(a::AbstractArray{T,D}, b::T) where {T,D} = LazyElementwiseOperation{T,D,:-}(a,b)
140 Base.@propagate_inbounds *̃(a::AbstractArray{T,D}, b::Real) where {T,D} = LazyElementwiseOperation{T,D,:*}(a,b) 70 Base.@propagate_inbounds *̃(a::AbstractArray{T,D}, b::T) where {T,D} = LazyElementwiseOperation{T,D,:*}(a,b)
141 Base.@propagate_inbounds /̃(a::AbstractArray{T,D}, b::Real) where {T,D} = LazyElementwiseOperation{T,D,:/}(a,b) 71 Base.@propagate_inbounds /̃(a::AbstractArray{T,D}, b::T) where {T,D} = LazyElementwiseOperation{T,D,:/}(a,b)
142 72
143 Base.@propagate_inbounds +̃(a::Real, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:+}(a,b) 73 Base.@propagate_inbounds +̃(a::T, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:+}(a,b)
144 Base.@propagate_inbounds -̃(a::Real, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:-}(a,b) 74 Base.@propagate_inbounds -̃(a::T, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:-}(a,b)
145 Base.@propagate_inbounds *̃(a::Real, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:*}(a,b) 75 Base.@propagate_inbounds *̃(a::T, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:*}(a,b)
146 Base.@propagate_inbounds /̃(a::Real, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:/}(a,b) 76 Base.@propagate_inbounds /̃(a::T, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:/}(a,b)
147 77
148 78
149 79
150 # NOTE: Är det knas att vi har till exempel * istället för .* ?? 80 # NOTE: Är det knas att vi har till exempel * istället för .* ??
151 # Oklart om det ens går att lösa.. 81 # Oklart om det ens går att lösa..