Type-stable broadcasting involving nested array

Hi Folks,

I am trying to get the following toy snippet type-stable, but wasn’t sure how to proceed.

using Random
using InteractiveUtils

using LinearAlgebra

const T = Float64


# scalar ver
function foo(
  a ::D,
  b ::AbstractVector{T};
  c ::AbstractVector{T}=similar(b),
) where {D<:Real, T<:Real}
  c[begin] = b[begin] + a
  return last(c)
end

# vector ver
function vfoo(
  va ::AbstractArray{D},
  vb ::AbstractArray{<:AbstractVector{T}};
  vc ::AbstractArray{<:AbstractVector{T}}=similar(b),
) where {D<:Real, T<:Real}

  ((a, b, c)->foo(a, b; c=c)).(va, vb, vc)

  return last.(vc)
end

# scalars
a = T(randn())
b = [T(randn())]
c = [T(randn())]

@code_warntype foo(a, b; c=c)

# vectors
va = [copy(a), copy(a)]
vb = [copy(b), copy(b)]
vc = [copy(c), copy(c)]

@code_warntype vfoo(va, vb; vc=vc)

The second @code_warntype complains about:

::Type{AbstractArray{var"#s6"<:AbstractVector{Float64}}}

and I don’t know how to deal with it.

Thanks ahead for your time and help!

BTW, links recommending tutorials of understanding @code_warntype outputs would also be very much appreciated!

if you remove all of your type annotations it’s actually stable

you didn’t copy the full line:

│   %8  = Core.apply_type(Main.AbstractArray, @_6::Core.Compiler.PartialTypeVar(var"#s6"<:AbstractVector{Float64}, true, true))::Type{AbstractArray{var"#s6"<:AbstractVector{Float64}}}

you see apply_type? your annotations are doing de-optimization

thanks for looking into this!

I thought it was a good practice to be specific on the intended usage of the function.

Is there a correct way to type-annotate vfoo?

maybe but why? annotating functions don’t speed them up, and only limit their applicability. if you’re mostly a user, you shouldn’t annotate too much

I was actually trying to make a package, and hoped to limit the API.

Was imagining users would be less confused if (type-stability) error shows at the top level they call, instead of during the broadcasting.

Have you tried this?

function vfoo(
    va::AbstractArray{D},
    vb::AbstractArray{A}, ;
    vc::AbstractArray{B}=similar(b),
) where {D<:Real, A<:AbstractVector{T}, B<:AbstractVector{T}} where T<:Real

    ((a, b, c)->foo(a, b; c)).(va, vb, vc)
    return last.(vc)
end

Why would you want to limit your API ? This seems like a shame : if someone want to use your code for something you did not know it could be used for, then they can’t.

I need to write an rrule for my API.

Awesome!

Would you also have time to share about how this one was different from the original one?
The type annotation of julia has been quite confusing for me, espacially those involving Abstract types…