Keys can be lost in `mapreduce(..., cat, ...)`

I want to slice a 3D array into matrices, multiply each matrix by another matrix, and then cat the result back into a 3D array.

The reason to do it this way is: the dimension shared by the matrices has non-overlapping keys, and I want to find the overlapping keys (that have non-missing values) for each slice.

The problem is that I can't preserve the axiskeys that I sliced over in the final concatenated array - it defaults to OneTo.

MWE:

julia> KA1 = KeyedArray(ones(2, 3), w=['a', 'b'], x=[:a, :b, :c]);

julia> KA2 = KeyedArray(ones(3, 4, 2), x=[:a, :b, :c], y=0.:3., z=["foo", "bar"]);

julia> mapreduce((x, y) -> cat(x, y; dims=:z), axiskeys(KA2, :z)) do z
           KA2_slice = KA2(z=z)
           return KA1 * KA2_slice
       end
3-dimensional KeyedArray(NamedDimsArray(...)) with keys:
↓   w ∈ 2-element Vector{Char}
→   y ∈ 4-element StepRangeLen{Float64,...}
□   z ∈ 2-element OneTo{Int}
And data, 2×4×2 Array{Float64, 3}:
[:, :, 1] ~ (:, :, 1):
         (0.0)  (1.0)  (2.0)  (3.0)
  ('a')    3.0    3.0    3.0    3.0
  ('b')    3.0    3.0    3.0    3.0

[:, :, 2] ~ (:, :, 2):
         (0.0)  (1.0)  (2.0)  (3.0)
  ('a')    3.0    3.0    3.0    3.0
  ('b')    3.0    3.0    3.0    3.0

I tried using an Interval at the KA2_slice = KA2(z=z) step, but Julia doesn't seem to support multiplying tensors with different dimensions, even if it's a trailing singleton dimension:

julia> ones(2, 3) * ones(3, 2, 1)
ERROR: MethodError: no method matching *(::Matrix{Float64}, ::Array{Float64, 3})
Closest candidates are:
  *(::Any, ::Any, ::Any, ::Any...) at operators.jl:560
  *(::StridedMatrix{T}, ::StridedVector{S}) where {T<:Union{Float32, Float64, ComplexF32, ComplexF64}, S<:Real} at /Users/julia/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/matmul.jl:44
  *(::StridedMatrix{var"#s832"} where var"#s832"<:Union{Float32, Float64}, ::StridedMatrix{var"#s831"} where var"#s831"<:Union{Float32, Float64, ComplexF32, ComplexF64}) at /Users/julia/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/matmul.jl:158
  ...
Stacktrace:
 [1] top-level scope
   @ REPL[21]:1

My current workaround is to use wrapdims afterward to re-key the array.

Start by reproducing the MWE involving mapreduce, cat, axiskeys(KA2, :z), and slicing with KA2(z=z). Trace how the z axis keys are handled during concatenation, then verify that the final 3D KeyedArray preserves the original z keys instead of defaulting to OneTo, without relying on wrapdims afterward.