Return type of `sum` with nested function is inferred as `Any`

General Usage Performance
1

I have a function that sums over expressions calling nested functions. I noticed that the function runs rather slowly. Upon further investigation, julia seems to be unable to infer the return type of this function.

The function looks somewhat like this:

struct Foo{T}
    a::T
    b::Array{T,1}
end

function bar(f::Foo)
    ab = axes(f.b, 1)
    baz(x) = sum(x * b for b in ab)
    sum(b * baz(f.a) for b in ab)
end

The call to sum in bar returns ::Any, according to @warn_type:

> f = Foo(3, [1, 2, 3])
Foo{Int64}(3, [1, 2, 3])
> @code_warntype bar(f)
Variables
  #self#::Core.Compiler.Const(bar, false)
  f::foo{Int64}
  #2::var"#2#5"{foo{Int64},var"#baz#3"{Base.OneTo{Int64}}}
  ab::Base.OneTo{Int64}
  baz::var"#baz#3"{Base.OneTo{Int64}}

Body::Any
1 ─ %1  = Base.getproperty(f, :b)::Array{Int64,1}
│         (ab = Main.axes(%1, 1))
│   %3  = Main.:(var"#baz#3")::Core.Compiler.Const(var"#baz#3", false)
│   %4  = Core.typeof(ab)::Core.Compiler.Const(Base.OneTo{Int64}, false)
│   %5  = Core.apply_type(%3, %4)::Core.Compiler.Const(var"#baz#3"{Base.OneTo{Int64}}, false)
│         (baz = %new(%5, ab))
│   %7  = Main.:(var"#2#5")::Core.Compiler.Const(var"#2#5", false)
│   %8  = Core.typeof(f)::Core.Compiler.Const(Foo{Int64}, false)
│   %9  = Core.typeof(baz)::Core.Compiler.Const(var"#baz#3"{Base.OneTo{Int64}}, false)
│   %10 = Core.apply_type(%7, %8, %9)::Core.Compiler.Const(var"#2#5"{Foo{Int64},var"#baz#3"{Base.OneTo{Int64}}}, false)
│         (#2 = %new(%10, f, baz))
│   %12 = #2::var"#2#5"{Foo{Int64},var"#baz#3"{Base.OneTo{Int64}}}
│   %13 = Base.Generator(%12, ab)::Base.Generator{Base.OneTo{Int64},var"#2#5"{Foo{Int64},var"#baz#3"{Base.OneTo{Int64}}}}
│   %14 = Main.sum(%13)::Any
└──       return %14

How can I debug this further and how can I help julia inferring this type?

I know of #15276, but I don’t believe this is applicable here. As far as I can tell, the variables from bar’s scope are not boxed.

I also tried Cthulhu.jl’s @descend macro, but wasn’t able to recognize the exact point, where the information is lost.

The function is inside a module, but I observe the same behavior, when I define struct and function in the REPL.

1 Like
2

This is interesting, looks like a bug in Generator? It would be really interesting to get to the bottom of it. Maybe you can open an issue?

For now, as a workaround, you can change generator to an anonymous functions

function bar(f::Foo)
    ab = axes(f.b, 1)
    baz(x) = sum(x * b for b in ab)
    sum(b -> b * baz(f.a), ab)
end
3

Thanks! I’ll try the lambda.
I’ve opened a bug, btw.

4

This works partly. For the example I gave in the first post, the sum+iterator solution appears to work as expected.

But my actual example has a sum(b * c * baz(f.a) for b in ab, c in ac) type of call. I’ve replaced this with sum((b,c) -> b * c * baz(f.a), Iterators.product(ab, ac)). Now instead of Any, the function body resolves to Union{}:

Variables
  #self#::Core.Compiler.Const(bar, false)
  f::foo{Int64}
  #7::var"#7#10"{foo{Int64},var"#baz#8"{Base.OneTo{Int64}}}
  ab::Base.OneTo{Int64}
  ac::Base.OneTo{Int64}
  baz::var"#baz#8"{Base.OneTo{Int64}}

Body::Union{}
1 ─ %1  = Base.getproperty(f, :b)::Array{Int64,1}
│         (ab = Main.axes(%1, 1))
│   %3  = Base.getproperty(f, :b)::Array{Int64,1}
│         (ac = Main.axes(%3, 2))
│   %5  = Main.:(var"#baz#8")::Core.Compiler.Const(var"#baz#8", false)
│   %6  = Core.typeof(ab)::Core.Compiler.Const(Base.OneTo{Int64}, false)
│   %7  = Core.apply_type(%5, %6)::Core.Compiler.Const(var"#baz#8"{Base.OneTo{Int64}}, false)
│         (baz = %new(%7, ab))
│   %9  = Main.:(var"#7#10")::Core.Compiler.Const(var"#7#10", false)
│   %10 = Core.typeof(f)::Core.Compiler.Const(foo{Int64}, false)
│   %11 = Core.typeof(baz)::Core.Compiler.Const(var"#baz#8"{Base.OneTo{Int64}}, false)
│   %12 = Core.apply_type(%9, %10, %11)::Core.Compiler.Const(var"#7#10"{foo{Int64},var"#baz#8"{Base.OneTo{Int64}}}, false)
│         (#7 = %new(%12, f, baz))
│   %14 = #7::var"#7#10"{foo{Int64},var"#baz#8"{Base.OneTo{Int64}}}
│   %15 = Base.getproperty(Main.Iterators, :product)::Core.Compiler.Const(Base.Iterators.product, false)
│   %16 = ab::Base.OneTo{Int64}
│   %17 = (%15)(%16, ac::Core.Compiler.Const(Base.OneTo(1), false))::Core.Compiler.PartialStruct(Base.Iterators.ProductIterator{Tuple{Base.OneTo{Int64},Base.
OneTo{Int64}}}, Any[Core.Compiler.PartialStruct(Tuple{Base.OneTo{Int64},Base.OneTo{Int64}}, Any[Base.OneTo{Int64}, Core.Compiler.Const(Base.OneTo(1), false)]
)])
│         Main.sum(%14, %17)
└──       Core.Compiler.Const(:(return %18), false)
5

But I believe this is actually because I can’t call the (b,c) -> ... lambda with a Tuple (b,c).

6

Well, it’s hard to solve without full example, but there are multiple ways to overcome this problem.

  1. It looks like baz(f.a) is not dependent on b. Then you can just calculate it outside the sum (and I would also remove baz from inside the bar, in 99% of cases there is no need to define internal functions)
function baz(x, ab)
    sum(x * b for b in ab)
end

function bar(f::Foo)
    ab = axes(f.b, 1)
    c = baz(f.a)
    sum(b * c for b in ab)
end

and the same can be done in multi loop case.

  1. If calculations cannot be moved outside the sum function, you can always use ordinary loops. It’s more verbose, but should have the same performance, without inference issues
function bar(f::Foo)
    ab = axes(f.b, 1)
    baz(x) = sum(x * b for b in ab)
    res = 0
    @inbounds for b in ab
        res += b * baz(f.a)
    end
    return res
end
7

The full example is a bit more complex, so I tried to give a MWE… Went a bit too minimal apparently :wink:

I have similar inference problems when using Iterators.product() properly (sum(t -> t[1] * t[2] * baz(f.a), Iterators.product(ab1, ab2))), so I guess I’ll go with the manual route.

Thanks for your help!

1 Like