diff src/LazyTensors/lazy_array.jl @ 1207:f1c2a4fa0ee1 performance/get_region_type_inference

Merge default
author Jonatan Werpers <jonatan@werpers.com>
date Fri, 03 Feb 2023 22:14:47 +0100
parents 6fd97b5ed3e5
children c5cf32fab163
line wrap: on
line diff
--- a/src/LazyTensors/lazy_array.jl	Mon Feb 21 10:33:58 2022 +0100
+++ b/src/LazyTensors/lazy_array.jl	Fri Feb 03 22:14:47 2023 +0100
@@ -28,7 +28,7 @@
 export LazyFunctionArray
 
 function LazyFunctionArray(f::F, size::NTuple{D,Int}) where {F<:Function,D}
-    T = typeof(f(ones(D)...))
+    T = typeof(f(ones(Int, D)...))
     return LazyFunctionArray{F,T,D}(f,size)
 end
 
@@ -36,7 +36,7 @@
 
 function Base.getindex(lfa::LazyFunctionArray{F,T,D}, I::Vararg{Int,D}) where {F,T,D}
     @boundscheck checkbounds(lfa, I...)
-    return lfa.f(I...)
+    return @inbounds lfa.f(I...)
 end
 
 
@@ -61,25 +61,17 @@
 end
 LazyElementwiseOperation{T,D,Op}(a::AbstractArray{T,D},b::T) where {T,D,Op} = LazyElementwiseOperation{T,D,Op}(a, LazyConstantArray(b, size(a)))
 LazyElementwiseOperation{T,D,Op}(a::T,b::AbstractArray{T,D}) where {T,D,Op} = LazyElementwiseOperation{T,D,Op}(LazyConstantArray(a,  size(b)), b)
-# 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...]
+Base.@propagate_inbounds evaluate(leo::LazyElementwiseOperation{T,D,:+}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] + leo.b[I...]
+Base.@propagate_inbounds evaluate(leo::LazyElementwiseOperation{T,D,:-}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] - leo.b[I...]
+Base.@propagate_inbounds evaluate(leo::LazyElementwiseOperation{T,D,:*}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] * leo.b[I...]
+Base.@propagate_inbounds 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}, I::Vararg{Int,D}) where {T,D}
-    @boundscheck if !checkbounds(Bool, leo.a, I...)
-        throw(BoundsError([leo], I...))
-    end
-    return evaluate(leo, I...)
+function Base.getindex(leo::LazyElementwiseOperation{T,D}, I::Vararg{Int,D}) where {T,D}
+    @boundscheck checkbounds(leo.a, I...)
+    return @inbounds evaluate(leo, I...)
 end
 
 # Define lazy operations for AbstractArrays. Operations constructs a LazyElementwiseOperation which
@@ -101,17 +93,22 @@
 
 
 
-# NOTE: Är det knas att vi har till exempel * istället för .* ??
-# Oklart om det ens går att lösa..
+# Overload +,-,*,/ for LazyArrays 
+# Element wise operation for `*` and `/` are not overloaded for due to conflicts with the behavior
+# of regular `*` and `/` for AbstractArrays. Use tilde versions instead.
 Base.@propagate_inbounds Base.:+(a::LazyArray{T,D}, b::LazyArray{T,D}) where {T,D} = a +̃ b
+Base.@propagate_inbounds Base.:-(a::LazyArray{T,D}, b::LazyArray{T,D}) where {T,D} = a -̃ b
+
 Base.@propagate_inbounds Base.:+(a::LazyArray{T,D}, b::AbstractArray{T,D}) where {T,D} = a +̃ b
-Base.@propagate_inbounds Base.:+(a::AbstractArray{T,D}, b::LazyArray{T,D}) where {T,D} = a +̃ b
+Base.@propagate_inbounds Base.:-(a::LazyArray{T,D}, b::AbstractArray{T,D}) where {T,D} = a -̃ b
 
-Base.@propagate_inbounds Base.:-(a::LazyArray{T,D}, b::LazyArray{T,D}) where {T,D} = a -̃ b
-Base.@propagate_inbounds Base.:-(a::LazyArray{T,D}, b::AbstractArray{T,D}) where {T,D} = a -̃ b
+Base.@propagate_inbounds Base.:+(a::AbstractArray{T,D}, b::LazyArray{T,D}) where {T,D} = a +̃ b
 Base.@propagate_inbounds Base.:-(a::AbstractArray{T,D}, b::LazyArray{T,D}) where {T,D} = a -̃ b
 
-# Element wise operation for `*` and `\` are not overloaded due to conflicts with the behavior
-# of regular `*` and `/` for AbstractArrays. Use tilde versions instead.
+Base.@propagate_inbounds Base.:+(a::LazyArray{T,D}, b::T) where {T,D} = a +̃ b
+Base.@propagate_inbounds Base.:-(a::LazyArray{T,D}, b::T) where {T,D} = a -̃ b
+
+Base.@propagate_inbounds Base.:+(a::T, b::LazyArray{T,D}) where {T,D} = a +̃ b
+Base.@propagate_inbounds Base.:-(a::T, b::LazyArray{T,D}) where {T,D} = a -̃  b
 
 export +̃, -̃, *̃, /̃