# Type-instability of mapreduce vs. map + reduce

**URL:** <https://discourse.julialang.org/t/type-instability-of-mapreduce-vs-map-reduce/121136>\
**Category:** Performance\
**Tags:** cat, jet, stack, mapreduce\
**Created:** [October 10, 2024, 10:26am UTC](https://discourse.julialang.org/t/type-instability-of-mapreduce-vs-map-reduce/121136 "2024-10-10T10:26:42Z")\
**Posts on this page:** 19\
**Page:** 1

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [October 10, 2024, 10:26am UTC](https://discourse.julialang.org/t/type-instability-of-mapreduce-vs-map-reduce/121136/1 "2024-10-10T10:26:43Z")

</div>

I’m a bit surprised that in this simple example, `f` is type-stable but `g` is not. Is this a JET.jl issue (@aviatesk) or a Julia issue?

```julia
using JET

function f(m, n)
    blocks = map(1:n) do i
        ones(m)
    end
    return reduce(hcat, blocks)
end

function g(m, n)
    return mapreduce(hcat, 1:n) do i
        ones(m)
    end
end 

```

```julia
julia> @assert f(10, 20) == g(10, 20)

julia> @test_opt f(10, 20)
Test Passed

julia> @test_opt g(10, 20)
JET-test failed at /home/guillaume/Work/GitHub/Julia/DifferentiationInterface.jl/DifferentiationInterfaceTest/test/playground.jl:18
  Expression: #= /home/guillaume/Work/GitHub/Julia/DifferentiationInterface.jl/DifferentiationInterfaceTest/test/playground.jl:18 =# JET.@test_opt g(10, 20)
  ═════ 1 possible error found ═════
  ┌ g(m::Int64, n::Int64) @ Main /home/guillaume/Work/GitHub/Julia/DifferentiationInterface.jl/DifferentiationInterfaceTest/test/playground.jl:11
  │┌ mapreduce(f::var"#23#24"{Int64}, op::typeof(hcat), A::UnitRange{Int64}) @ Base ./reducedim.jl:357
  ││┌ mapreduce(f::var"#23#24"{Int64}, op::typeof(hcat), A::UnitRange{Int64}; dims::Colon, init::Base._InitialValue) @ Base ./reducedim.jl:357
  │││┌ _mapreduce_dim(f::var"#23#24"{Int64}, op::typeof(hcat), ::Base._InitialValue, A::UnitRange{Int64}, ::Colon) @ Base ./reducedim.jl:365
  ││││┌ _mapreduce(f::var"#23#24"{Int64}, op::typeof(hcat), ::IndexLinear, A::UnitRange{Int64}) @ Base ./reduce.jl:432
  │││││┌ mapreduce_empty_iter(f::var"#23#24"{Int64}, op::typeof(hcat), itr::UnitRange{Int64}, ItrEltype::Base.HasEltype) @ Base ./reduce.jl:380
  ││││││┌ reduce_empty_iter(op::Base.MappingRF{var"#23#24"{Int64}, typeof(hcat)}, itr::UnitRange{Int64}, ::Base.HasEltype) @ Base ./reduce.jl:384
  │││││││┌ reduce_empty(op::Base.MappingRF{var"#23#24"{Int64}, typeof(hcat)}, ::Type{Int64}) @ Base ./reduce.jl:361
  ││││││││ runtime dispatch detected: Base.mapreduce_empty(%1::var"#23#24"{Int64}, hcat, ::Int64)
  │││││││└────────────────────
  
ERROR: There was an error during testing

```

@Oscar_Smith and @Mason, sorry for the ping, @hill tells me you may have intuitions about this.

---

<div class="post-metadata">

**Author:** ![Mason](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mason/32/2423_2.png) [@Mason](https://discourse.julialang.org/u/Mason)\
**Post date:** [October 10, 2024, 10:52am UTC](https://discourse.julialang.org/t/type-instability-of-mapreduce-vs-map-reduce/121136/2 "2024-10-10T10:52:21Z")

</div>

Yeah, this is because we have a special codepath for `reduce(hcat, arr)` but not for `mapreduce(f, hcat, arr)`.

The special overloads for `reduce(*cat)` are really brittle and easily broken, which is why @mcabbott took on creating `stack` which is what I’d recommend using instead.

```julia
h(m, n) = stack(1:n) do i
    ones(m)
end

```

```julia-repl
julia> h(10, 20) == g(10, 20)
true

julia> @test_opt h(10, 20)
Test Passed

```

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [October 10, 2024, 10:58am UTC](https://discourse.julialang.org/t/type-instability-of-mapreduce-vs-map-reduce/121136/3 "2024-10-10T10:58:11Z")

</div>

Thanks for the quick answer! I’m aware of `stack` but it doesn’t fit my purpose for two reasons:

First, the actual use case involves concatenating smaller matrices into a bigger matrix, without adding an extra dimension. I could `stack` into an `Array{_,3}` and then drop the last dimension, but that seems wasteful?

```julia
function g(m, n)
    return mapreduce(hcat, 1:n) do i
        ones(m, 2) # this is a block and not just a column
    end
end 

```

Second, `stack` doesn’t have an optimized implementation for static arrays, and it actually returns a `Matrix` in those cases. I would be willing to try and fix it, but I’d need some help, maybe from @mcabbott.

> <https://github.com/JuliaArrays/StaticArrays.jl/issues/1272>
>
> x/ref https://github.com/EnzymeAD/Enzyme.jl/issues/1714

Is there another way out of this?

---

<div class="post-metadata">

**Author:** ![Mason](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mason/32/2423_2.png) [@Mason](https://discourse.julialang.org/u/Mason)\
**Post date:** [October 10, 2024, 11:03am UTC](https://discourse.julialang.org/t/type-instability-of-mapreduce-vs-map-reduce/121136/4 "2024-10-10T11:03:36Z")

</div>

I’d probably just write a `for` loop in that case, or use Tullio.jl, I think it knows about staticarrays (but I’m not sure).

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [October 10, 2024, 11:04am UTC](https://discourse.julialang.org/t/type-instability-of-mapreduce-vs-map-reduce/121136/5 "2024-10-10T11:04:50Z")

</div>

Unfortunately this is for DifferentiationInterface.jl out-of-place Jacobian matrices, so

- I can’t pre-allocate and then mutate in a `for` loop, for instance because it would waste the benefits of static arrays
- I can’t pull in something like Tullio.jl because the dependencies need to be absolutely minimal

---

<div class="post-metadata">

**Author:** ![Mason](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mason/32/2423_2.png) [@Mason](https://discourse.julialang.org/u/Mason)\
**Post date:** [October 10, 2024, 11:23am UTC](https://discourse.julialang.org/t/type-instability-of-mapreduce-vs-map-reduce/121136/6 "2024-10-10T11:23:47Z")

</div>

> [@gdalle](#):
>
> I can’t pre-allocate and then mutate in a `for` loop, for instance because it would waste the benefits of static arrays

Yes you can. Just allocate a `MArray`, mutate it, then convert to `SArray`.

```julia-repl
julia> using StaticArrays

julia> function f(::Val{m}, ::Val{n}) where {m, n}
           M = MArray{Tuple{m, n}, Float64}(undef)
           for j ∈ axes(m, 2)
               for i ∈ axes(m, 1)
                   M[i, j] = 1
               end
           end
           SArray(M)
       end
f (generic function with 1 method)

julia> @btime f(Val(10), Val(20))
  42.968 ns (0 allocations: 0 bytes)
10×20 SMatrix{10, 20, Float64, 200} with indices SOneTo(10)×SOneTo(20):
 1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0

```

The trick here is that you can create `MArray` objects without any allocations **if** the compiler knows that the array never escapes the function body. That’s why we convert it to an `SArray` at the end.

In v1.12, this will get better becuase the compiler can reason about non-inlined function calls, but even in quite early versions of julia, simple kernels like this where everything is inlined works really well.

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [October 10, 2024, 11:44am UTC](https://discourse.julialang.org/t/type-instability-of-mapreduce-vs-map-reduce/121136/7 "2024-10-10T11:44:25Z")

</div>

Good to know that this works! The challenge now will be to code it in a way that is not specialized on StaticArrays but instead works the same on most `AbstractArray`s.  
And I still think optimizing `stack` on `SArray`s would be worth doing, if anyone has pointers.

---

<div class="post-metadata">

**Author:** ![mbauman](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mbauman/32/31082_2.png) [@mbauman](https://discourse.julialang.org/u/mbauman)\
**Post date:** [October 10, 2024, 12:40pm UTC](https://discourse.julialang.org/t/type-instability-of-mapreduce-vs-map-reduce/121136/8 "2024-10-10T12:40:19Z")

</div>

`g` is type unstable because the reducer — `hcat` — is not called for one-value collections:

```Julia
julia> g(3, 1)
3-element Vector{Float64}:
 1.0
 1.0
 1.0

julia> g(3, 2)
3×2 Matrix{Float64}:
 1.0 1.0
 1.0 1.0
 1.0 1.0

```

To ensure type stability, use `ones(m, 1)`.

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [October 10, 2024, 12:52pm UTC](https://discourse.julialang.org/t/type-instability-of-mapreduce-vs-map-reduce/121136/9 "2024-10-10T12:52:52Z")

</div>

That applies to the first example I gave, but apparently not to my real use case with blocks, since `hcat`-ing even just one always returns a matrix?

```julia
julia> using JET

julia> function g(m, n)
           return mapreduce(hcat, 1:n) do i
               ones(m, 2)
           end
       end
g (generic function with 1 method)

julia> g(3, 1)
3×2 Matrix{Float64}:
 1.0 1.0
 1.0 1.0
 1.0 1.0

julia> @test_opt g(3, 1)
JET-test failed at /home/guillaume/Work/GitHub/Julia/DifferentiationInterface.jl/DifferentiationInterface/test/playground.jl:11
  Expression: #= /home/guillaume/Work/GitHub/Julia/DifferentiationInterface.jl/DifferentiationInterface/test/playground.jl:11 =# JET.@test_opt g(3, 1)
  ═════ 1 possible error found ═════
  ┌ g(m::Int64, n::Int64) @ Main /home/guillaume/Work/GitHub/Julia/DifferentiationInterface.jl/DifferentiationInterface/test/playground.jl:4
  │┌ mapreduce(f::var"#35#36"{Int64}, op::typeof(hcat), A::Base.OneTo{Int64}) @ Base ./reducedim.jl:357
  ││┌ mapreduce(f::var"#35#36"{Int64}, op::typeof(hcat), A::Base.OneTo{Int64}; dims::Colon, init::Base._InitialValue) @ Base ./reducedim.jl:357
  │││┌ _mapreduce_dim(f::var"#35#36"{Int64}, op::typeof(hcat), ::Base._InitialValue, A::Base.OneTo{Int64}, ::Colon) @ Base ./reducedim.jl:365
  ││││┌ _mapreduce(f::var"#35#36"{Int64}, op::typeof(hcat), ::IndexLinear, A::Base.OneTo{Int64}) @ Base ./reduce.jl:432
  │││││┌ mapreduce_empty_iter(f::var"#35#36"{Int64}, op::typeof(hcat), itr::Base.OneTo{Int64}, ItrEltype::Base.HasEltype) @ Base ./reduce.jl:380
  ││││││┌ reduce_empty_iter(op::Base.MappingRF{var"#35#36"{Int64}, typeof(hcat)}, itr::Base.OneTo{Int64}, ::Base.HasEltype) @ Base ./reduce.jl:384
  │││││││┌ reduce_empty(op::Base.MappingRF{var"#35#36"{Int64}, typeof(hcat)}, ::Type{Int64}) @ Base ./reduce.jl:361
  ││││││││ runtime dispatch detected: Base.mapreduce_empty(%1::var"#35#36"{Int64}, hcat, ::Int64)
  │││││││└────────────────────
  
ERROR: There was an error during testing

```

---

<div class="post-metadata">

**Author:** ![mbauman](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mbauman/32/31082_2.png) [@mbauman](https://discourse.julialang.org/u/mbauman)\
**Post date:** [October 10, 2024, 1:01pm UTC](https://discourse.julialang.org/t/type-instability-of-mapreduce-vs-map-reduce/121136/10 "2024-10-10T13:01:13Z")

</div>

That’s in the error path for empty collections. You can avoid that path if you can use an `init`: `init=Array{Float64}(undef, m, 0)`. And using an `init` will also ensure that `hcat` is called for 1-arg collections, too.

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [October 10, 2024, 1:10pm UTC](https://discourse.julialang.org/t/type-instability-of-mapreduce-vs-map-reduce/121136/11 "2024-10-10T13:10:33Z")

</div>

I think that’s the trick I was missing! Thank you.

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [October 10, 2024, 1:57pm UTC](https://discourse.julialang.org/t/type-instability-of-mapreduce-vs-map-reduce/121136/12 "2024-10-10T13:57:45Z")

</div>

Now of course I need to figure out a function which generates a zero-width equivalent for any `AbstractMatrix`, to use in the `init`…

---

<div class="post-metadata">

**Author:** ![mbauman](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mbauman/32/31082_2.png) [@mbauman](https://discourse.julialang.org/u/mbauman)\
**Post date:** [October 10, 2024, 2:12pm UTC](https://discourse.julialang.org/t/type-instability-of-mapreduce-vs-map-reduce/121136/13 "2024-10-10T14:12:18Z")

</div>

Are you chasing down dynamic dispatches for static compilation? If not, you really don’t need to worry about this.

But if you are chasing every single dynamic dispatch down, another alternative could be to locally patch in Julia#51948 — prior to that we were dynamically fixing up a method error.

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [October 10, 2024, 2:34pm UTC](https://discourse.julialang.org/t/type-instability-of-mapreduce-vs-map-reduce/121136/14 "2024-10-10T14:34:19Z")

</div>

> [@mbauman](#):
>
> Are you chasing down dynamic dispatches for static compilation? If not, you really don’t need to worry about this.

Yeah, this whole quest is for DifferentiationInterface.jl to perform optimally on StaticArrays.jl. In that case, I would ideally avoid any unnecessary runtime dispatch and even allocations.

> [@mbauman](#):
>
> But if you are chasing every single dynamic dispatch down, another alternative could be to locally patch in [Julia#51948](https://github.com/JuliaLang/julia/issues/51948) — prior to that we were dynamically fixing up a method error.

I’m not sure I understand, what do you want me to patch there?

---

<div class="post-metadata">

**Author:** ![mbauman](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mbauman/32/31082_2.png) [@mbauman](https://discourse.julialang.org/u/mbauman)\
**Post date:** [October 10, 2024, 2:37pm UTC](https://discourse.julialang.org/t/type-instability-of-mapreduce-vs-map-reduce/121136/15 "2024-10-10T14:37:08Z")

</div>

The magic line is `@eval Base mapreduce_empty(f, op, T) = _empty_reduce_error()`. But probably not something you want to do in this context.

If you use a named function for the mapper, then you could overload this without piracy or shame.

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [October 10, 2024, 2:51pm UTC](https://discourse.julialang.org/t/type-instability-of-mapreduce-vs-map-reduce/121136/16 "2024-10-10T14:51:07Z")

</div>

I think I’m gonna go with another solution and compute the first matrix block outside of the `mapreduce`. Less elegant but less scary. Thanks a lot for your advice!

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [October 10, 2024, 4:35pm UTC](https://discourse.julialang.org/t/type-instability-of-mapreduce-vs-map-reduce/121136/17 "2024-10-10T16:35:09Z")

</div>

Do you think there is a way to tell the compiler statically that the length of the iterator won’t be zero?

---

<div class="post-metadata">

**Author:** ![mbauman](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mbauman/32/31082_2.png) [@mbauman](https://discourse.julialang.org/u/mbauman)\
**Post date:** [October 10, 2024, 4:40pm UTC](https://discourse.julialang.org/t/type-instability-of-mapreduce-vs-map-reduce/121136/18 "2024-10-10T16:40:11Z")

</div>

I’d take a step back — are you only chasing this because JET is telling you to chase it? I’m not sure it’s actually a real dynamic dispatch; it’s an _intentional_ method missing error. And it’s not there in v1.11.

You might be ok to ignore it, I think.

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [October 10, 2024, 4:43pm UTC](https://discourse.julialang.org/t/type-instability-of-mapreduce-vs-map-reduce/121136/19 "2024-10-10T16:43:03Z")

</div>

You’re probably right. JET has been bossing me around for too long.  
And I was testing on 1.10, so good catch. Let’s close this thread.