comparison src/LazyTensors/lazy_array.jl @ 1004:7fd37aab84fe refactor/lazy_tensors

Simplify bounds handling for LazyElementwiseOperation
author Jonatan Werpers <jonatan@werpers.com>
date Sun, 20 Mar 2022 21:35:20 +0100
parents 7ef605b8f132
children 4dd3c2312d9e
comparison
equal deleted inserted replaced
1003:7ef605b8f132 1004:7fd37aab84fe
62 LazyElementwiseOperation{T,D,Op}(a::AbstractArray{T,D},b::T) where {T,D,Op} = LazyElementwiseOperation{T,D,Op}(a, LazyConstantArray(b, size(a))) 62 LazyElementwiseOperation{T,D,Op}(a::AbstractArray{T,D},b::T) where {T,D,Op} = LazyElementwiseOperation{T,D,Op}(a, LazyConstantArray(b, size(a)))
63 LazyElementwiseOperation{T,D,Op}(a::T,b::AbstractArray{T,D}) where {T,D,Op} = LazyElementwiseOperation{T,D,Op}(LazyConstantArray(a, size(b)), b) 63 LazyElementwiseOperation{T,D,Op}(a::T,b::AbstractArray{T,D}) where {T,D,Op} = LazyElementwiseOperation{T,D,Op}(LazyConstantArray(a, size(b)), b)
64 64
65 Base.size(v::LazyElementwiseOperation) = size(v.a) 65 Base.size(v::LazyElementwiseOperation) = size(v.a)
66 66
67 evaluate(leo::LazyElementwiseOperation{T,D,:+}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] + leo.b[I...] 67 evaluate(leo::LazyElementwiseOperation{T,D,:+}, I::Vararg{Int,D}) where {T,D} = @inbounds leo.a[I...] + leo.b[I...]
68 evaluate(leo::LazyElementwiseOperation{T,D,:-}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] - leo.b[I...] 68 evaluate(leo::LazyElementwiseOperation{T,D,:-}, I::Vararg{Int,D}) where {T,D} = @inbounds leo.a[I...] - leo.b[I...]
69 evaluate(leo::LazyElementwiseOperation{T,D,:*}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] * leo.b[I...] 69 evaluate(leo::LazyElementwiseOperation{T,D,:*}, I::Vararg{Int,D}) where {T,D} = @inbounds leo.a[I...] * leo.b[I...]
70 evaluate(leo::LazyElementwiseOperation{T,D,:/}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] / leo.b[I...] 70 evaluate(leo::LazyElementwiseOperation{T,D,:/}, I::Vararg{Int,D}) where {T,D} = @inbounds leo.a[I...] / leo.b[I...]
71 71
72 Base.@propagate_inbounds @inline function Base.getindex(leo::LazyElementwiseOperation{T,D}, I::Vararg{Int,D}) where {T,D} 72 function Base.getindex(leo::LazyElementwiseOperation{T,D}, I::Vararg{Int,D}) where {T,D}
73 @boundscheck if !checkbounds(Bool, leo.a, I...) 73 @boundscheck checkbounds(leo.a, I...)
74 throw(BoundsError([leo], I...))
75 end
76 return evaluate(leo, I...) 74 return evaluate(leo, I...)
77 end 75 end
78 76
79 # Define lazy operations for AbstractArrays. Operations constructs a LazyElementwiseOperation which 77 # Define lazy operations for AbstractArrays. Operations constructs a LazyElementwiseOperation which
80 # can later be indexed into. Lazy operations are denoted by the usual operator followed by a tilde 78 # can later be indexed into. Lazy operations are denoted by the usual operator followed by a tilde