Nested `map` with a capturing closure is not inferred

General Usage Performance
, ,
1

A map whose function itself calls map with a closure capturing a variable loses its inferred element type, even though each map alone is inferred fine:

using Test

h(t, d) = map(u -> u + d, t)                  # inner map alone
f(ts, d) = map(t -> map(u -> u + d, t), ts)   # nested map, inner closure captures `d`
g(ts) = map(t -> map(u -> u + 1, t), ts)      # nested map, inner closure captures nothing

ts = [[1.0, 2.0], [3.0]]

@inferred h(ts[1], 1.0)
@inferred g(ts)
@inferred f(ts, 1.0)
1.10.12 1.11.9 1.12.7 1.13.1
@inferred h(ts[1], 1.0) passes passes passes passes
@inferred g(ts) fails passes passes passes
@inferred f(ts, 1.0) fails fails fails fails

The failure is:

return type Vector{Vector{Float64}} does not match inferred return type Vector

and @code_warntype f(ts, 1.0) shows Body::Vector. At runtime the result is a concrete Vector{Vector{Float64}}, so only the static return type is lost.

Surprisingly, with

h(t, d) = map(u -> u + d, t)
f2(ts, d) = map(t -> h(t, d), ts)

@inferred f2(ts, 1.0) fails in a fresh session, but passes if @inferred h(ts[1], 1.0) was run first. So the same code can pass or fail @inferred depending on what has already been compiled. (In f above this doesn’t help, because its inner closure is a different type from the one inside h.)

Swapping the outer map for broadcasting doesn’t fully avoid it either: (t -> map(x -> map(u -> u + d, x), t)).(tss) with tss = [[[1.0, 2.0]], [[3.0]]] also fails @inferred on 1.10 and 1.13.

2

As far as I can tell from recurring topics, type inference is intentionally terminated at recursive method calls with different input types in case it would meander endlessly without reaching a base case. It’s obvious to us that unrestricted type inference of f would terminate as g does, but I don’t know how or where the heuristic draws the line. I do know that input parametric types becoming less nested does not prove termination because any operation could make more complicated types, and type inference stopped before it could check that +(u::Float64, d::Float64)::Float64 does not call map(f, A::AbstractArray).

Return type inference, at least in reflection, is apparently stateful as you demonstrated with f2, but it’s worth noting that JET.@report_opt f2(ts, 1.0) reports recursion obstructing internal optimization either way as in @report_opt f(ts, 1.0).