Broadcasting over an array of custom struct gives type instability

General Usage
,
1

Hello,

I have a simple problem in Julia v1.10

using Test

struct MyMatrix{MT<:AbstractMatrix}
    M::MT
end

Base.:(*)(A::MyMatrix, x::Number) = MyMatrix(A.M * x)

function f(x)
    # return [y / 10 for y in x]
    return x * 10
end

# states = [rand_dm(10) for _ in 1:10]
states = [MyMatrix(rand(10, 10)) for _ in 1:10]

f(states)

@inferred f(states)

It detects type instability, since it is using broadcasting. But as soon as I use the comprehension, then everything works. Why?

Notice that the operation * itself is type-stable.

Notice also that the matrix can be any AbstractMatrix sparse, GPU, etc.

2

I’ll leave the analysis of the vector times number method to someone else but if you instead use an explicit broadcast this problem goes away:

f(x) = x .* 10
4

I cannot see that using a global const or a local scope makes any difference, neither on 1.10 nor on nightly. It also doesn’t make much sense that it would; the inference is made for the type present when the call is done. It’s something completely different to refer to a non-const global from inside a function, rather than a function argument, which is a common cause of type instability.

5

The indirect broadcast gets bad inference at the Base.broadcast_preserving_zero_d call, I don’t know how @code_warntype f(states) misses it:

julia> @inferred f(states)
ERROR: return type Vector{MyMatrix{Matrix{Float64}}} does not match inferred return type Any
Stacktrace:
 [1] error(s::String)
   @ Base .\error.jl:35
 [2] top-level scope
   @ REPL[68]:1

julia> @code_warntype f(states)
MethodInstance for f(::Vector{MyMatrix{Matrix{Float64}}})
  from f(x) @ Main REPL[50]:1
Arguments
  #self#::Core.Const(Main.f)
  x::Vector{MyMatrix{Matrix{Float64}}}
Body::Vector{MyMatrix{Matrix{Float64}}}
1 ─ %1 = Main.:*::Core.Const(*)
│   %2 = (%1)(x, 10)::Vector{MyMatrix{Matrix{Float64}}}
└──      return %2

julia> @code_warntype *(states, 10)
MethodInstance for *(::Vector{MyMatrix{Matrix{Float64}}}, ::Int64)
  from *(A::AbstractArray, B::Number) @ Base arraymath.jl:24
Arguments
  #self#::Core.Const(*)
  A::Vector{MyMatrix{Matrix{Float64}}}
  B::Int64
Body::Any
1 ─ %1 = Base.broadcast_preserving_zero_d::Core.Const(Base.Broadcast.broadcast_preserving_zero_d)
│   %2 = (%1)(Base.:*, A, B)::Any
└──      return %2
6

Thanks for correcting.
It seems the issue goes away by doing some things, which I wrongly attributed to the introduction of the local scope.
E.g. re-executing the struct definition makes this issue vanish.

julia> @inferred f(states);
ERROR: return type Vector{MyMatrix{Matrix{Float64}}} does not match inferred return type Any
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] top-level scope
   @ REPL[7]:1

julia> struct MyMatrix{MT<:AbstractMatrix}
           M::MT
       end

julia> @inferred f(states);
7

It also makes @code_warntype *(states, 10) type-stable and match @code_warntype f(states), what on earth?

Might be related to

because re-evaluating struct definitions with no changes currently keep the type but reevaluates the inner and default constructors.