JuliaArrays / BlockArrays.jl

BlockArrays for Julia
http://juliaarrays.github.io/BlockArrays.jl/
Other
198 stars 29 forks source link

Type instability of `blocksizes(y, i)` #425

Open lkdvos opened 2 weeks ago

lkdvos commented 2 weeks ago

I'm running into a subtle type-instability with blocksizes, which breaks the type-inference of a function I am trying to define further down the line.

using BlockArrays
x = BlockedArray(rand(4, 4), [2, 2], [2, 2])

@inferred blocksizes(x, 1)
ERROR: return type Vector{Int64} does not match inferred return type Union{FillArrays.Fill{Int64, 1, Tuple{Base.OneTo{Int64}}}, Vector{Int64}}

I did some digging, and it seems that the culprit is that Base.axes(::AbstractArray, n) will output a OneTo whenever n is larger than the number of dimensions. It seems like the compiler isn't able to const-propagate the n to figure this out. I was wondering if there is a way to avoid this instability, maybe by defining a fallback for axes(::AbstractBlockArray, d::Int) to output a BlockedOneTo instead?

I can definitely open a PR to make changes and tests, I just wanted to ask for some comments and feedback first.

jishnub commented 1 week ago

This should be fixed by https://github.com/JuliaArrays/BlockArrays.jl/pull/429, but this would work only within a function, and not at the top-level. This is because the dimension needs to be constant-propagated, and such optimizations won't happen at the top level.