JuliaArrays/LazyArrays.jl

LazyArray constructions in broadcasting calls fail if arrays are created in the call instead of referred to by their name

jishnub opened this issue · 8 comments

This works:

julia> a = ones(2,2);

julia> LazyArray(@~ @. a * 2)
(2×2 Array{Float64,2}) .* (Int64):
 2.0  2.0
 2.0  2.0

but attempting to create the array directly in the call doesn't

julia> LazyArray(@~ @. ones(2,2) * 2)
ERROR: MethodError: no method matching BroadcastArray{Float64,N,F,Args} where Args where F where N(::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{0},Nothing,typeof(*),Tuple{Array{Float64,2},Int64}})
Closest candidates are:
  BroadcastArray{Float64,N,F,Args} where Args where F where N(::Base.Broadcast.Broadcasted{var"#s22",var"#s21",var"#s20",var"#s17"} where var"#s17"<:Tuple where var"#s20" where var"#s21"<:Tuple{Vararg{Any,N}} where var"#s22"<:Union{Nothing, Base.Broadcast.BroadcastStyle}) where {T, N} at /home/jishnu/.julia/packages/LazyArrays/0Tiyn/src/lazybroadcasting.jl:39
Stacktrace:
 [1] _BroadcastArray(::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{0},Nothing,typeof(*),Tuple{Array{Float64,2},Int64}}) at /home/jishnu/.julia/packages/LazyArrays/0Tiyn/src/lazybroadcasting.jl:53
 [2] BroadcastArray(::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{0},Nothing,typeof(*),Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{0},Nothing,typeof(ones),Tuple{Int64,Int64}},Int64}}) at /home/jishnu/.julia/packages/LazyArrays/0Tiyn/src/lazybroadcasting.jl:54
 [3] LazyArray(::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{0},Nothing,typeof(*),Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{0},Nothing,typeof(ones),Tuple{Int64,Int64}},Int64}}) at /home/jishnu/.julia/packages/LazyArrays/0Tiyn/src/lazybroadcasting.jl:35
 [4] top-level scope at REPL[37]:1

Is this not something that is supported?

Note the two are slightly different, the latter is ones.(2,2) .* 2. But it still should have worked:

julia> ones.(2,2) .* 2
2×2 Array{Float64,2}:
 2.0  2.0
 2.0  2.0

Hey,
I found it somehow complicated. Why isn't it really possible to use ones within an @~ expression?

julia> @~ ones((10, 10))
Applied(ones,(10, 10))

julia> @~ (1 .+ ones((10, 10)))
ERROR: MethodError: no method matching size(::Applied{LazyArrays.DefaultApplyStyle, typeof(ones), Tuple{Tuple{Int64, Int64}}})
Closest candidates are:
  size(::Union{LinearAlgebra.QR, LinearAlgebra.QRCompactWY, LinearAlgebra.QRPivoted}) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/qr.jl:524
  size(::Union{LinearAlgebra.QR, LinearAlgebra.QRCompactWY, LinearAlgebra.QRPivoted}, ::Integer) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/qr.jl:523
  size(::Union{LinearAlgebra.Cholesky, LinearAlgebra.CholeskyPivoted}) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/cholesky.jl:442
  ...
Stacktrace:
 [1] axes
   @ ./abstractarray.jl:89 [inlined]
 [2] combine_axes
   @ ./broadcast.jl:484 [inlined]
 [3] instantiate(bc::Base.Broadcast.Broadcasted{Base.Broadcast.Unknown, Nothing, typeof(+), Tuple{Int64, Applied{LazyArrays.DefaultApplyStyle, typeof(ones), Tuple{Tuple{Int64, Int64}}}}})
   @ Base.Broadcast ./broadcast.jl:266
 [4] top-level scope
   @ REPL[97]:1

julia> @~ (1 .+ ones.((10, 10)))
Base.Broadcast.Broadcasted(+, (1, Base.Broadcast.Broadcasted(ones, ((10, 10),))))

julia> materialize(@~ (1 .+ ones.((10, 10))))
ERROR: MethodError: no method matching +(::Int64, ::Vector{Float64})
For element-wise addition, use broadcasting with dot syntax: scalar .+ array
Closest candidates are:
  +(::Any, ::Any, ::Any, ::Any...) at operators.jl:560
  +(::T, ::T) where T<:Union{Int128, Int16, Int32, Int64, Int8, UInt128, UInt16, UInt32, UInt64, UInt8} at int.jl:87
  +(::Union{Int16, Int32, Int64, Int8}, ::BigInt) at gmp.jl:534
  ...
Stacktrace:
 [1] _broadcast_getindex_evalf
   @ ./broadcast.jl:648 [inlined]
 [2] _broadcast_getindex
   @ ./broadcast.jl:621 [inlined]
 [3] (::Base.Broadcast.var"#19#20"{Base.Broadcast.Broadcasted{Base.Broadcast.Style{Tuple}, Nothing, typeof(+), Tuple{Int64, Base.Broadcast.Broadcasted{Base.Broadcast.Style{Tuple}, Nothing, typeof(ones), Tuple{Tuple{Int64, Int64}}}}}})(k::Int64)
   @ Base.Broadcast ./broadcast.jl:1098
 [4] ntuple
   @ ./ntuple.jl:49 [inlined]
 [5] copy
   @ ./broadcast.jl:1098 [inlined]
 [6] materialize(bc::Base.Broadcast.Broadcasted{Base.Broadcast.Style{Tuple}, Nothing, typeof(+), Tuple{Int64, Base.Broadcast.Broadcasted{Base.Broadcast.Style{Tuple}, Nothing, typeof(ones), Tuple{Tuple{Int64, Int64}}}}})
   @ Base.Broadcast ./broadcast.jl:883
 [7] top-level scope
   @ REPL[99]:1

A very simple example that shows the problem:

x = ones(10,10); 
q = LazyArray(@~ x .* abs(1.0))

It seems like any function call not operating on an array causes problems.

@tkf wrote the @~ macro, perhaps he has some idea how to fix this?

I think you are right.

materialize(applied(x->x.*abs(x),x))

works fine.

Er it's not an issue with @~. This is a fix:

julia> Base.size(a::Applied) = size(materialize(a))

julia> q = LazyArray(@~ x .* abs(1.0))
(10×10 Matrix{Float64}) .* (Applied{LazyArrays.DefaultApplyStyle, typeof(abs), Tuple{Float64}}):
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0

Anyone feel like preparing a PR with this change and tests?

Whilst materializing to determine the size in general defeats the point its fine for these scalar cases, and should not be called where the current code already works

Actually ideally we would want it to always materialize when its scalar... it might be possible to make LazyArray(::Applied) materialize all such calls but unfortunately I don't have time to look at this right now.

One way would be to be smart at the applied stage. Unfortunately broadcasted doesn't already do this for scalars, only for ranges and a few others:

julia> Base.broadcasted(+, 1, 2:3) # smart and materializes as non-allocated
3:4

julia> Base.broadcasted(+, 1, 2) # dumb, does not replace with 3
Base.Broadcast.Broadcasted(+, (1, 2))

But I think something like the following might be an approach:

# determine ApplyStyle at `applied`, not `Applied`:
applied(f, x...) = applied(combine_apply_style(f, args...), f, x...)
applied(style, f, x...) = Applied{typeof(style)}(f, x...)

struct MaterializeStyle <: ApplyStyle end
ApplyStyle(f, x::Type{<:Number}...) = Base.promoteop(f, x) isa Number ? MaterializeStyle() : DefaultApplyStyle()
applied(::MaterializeStyle, f, x...) = f(x...)

While promoteop is warned against, I think its fine here since it would just resort to DefaultApplyStyle()