Error in gradient of stacked vectors

[mcabbott]

trafficstars

This should work:

julia> using TensorCast, Zygote

julia> f1(x,y) = [x^2, x*y, y^2];

julia> x = 1:2; y = 1:4;

julia> @cast out[i,j,k] := f1(x[i], y[j])[k]
2×4×3 transmute(stack(::Matrix{Vector{Int64}}), (2, 3, 1)) with eltype Int64:
[:, :, 1] =
 1  1  1  1
 4  4  4  4

[:, :, 2] =
 1  2  3  4
 2  4  6  8

[:, :, 3] =
 1  4  9  16
 1  4  9  16

julia> gradient(x -> sum(@cast out[i,j,k] := f1(x[i], y[j])[k]), 1:2)
ERROR: BoundsError: attempt to access 3×2×4 transmute(::FillArrays.Fill{Int64, 3, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}, (3, 1, 2)) with eltype Int64 at index [1:3, 1]