Here’s a minimal example of an error I’ve encountered with my code using Flux.

```
const x = zeros(Float32, 10, 4)
p = Flux.params(x)
function f()
x1 = softmax(x, dims=2)
x_d = [Diagonal(x1[:, k]) for k = 1:4]
x_d = cat(x_d..., dims=3)
# (10,10,4) array with dims=(1,2) diagonalized
return sum(x_d.^2)
end
g = gradient(p) do
loss = f()
end
```

While the forward pass runs to the function `f()`

runs without any problems, the gradient computations seems to throw this kind of an error:

```
ERROR: MethodError: no method matching _Diagonal_pullback(::Array{Float32, 3})
Closest candidates are:
_Diagonal_pullback(::Diagonal) at ~/.julia/packages/ChainRules/3yDBX/src/rulesets/LinearAlgebra/structured.jl:54
_Diagonal_pullback(::AbstractMatrix) at ~/.julia/packages/ChainRules/3yDBX/src/rulesets/LinearAlgebra/structured.jl:53
_Diagonal_pullback(::ChainRulesCore.Tangent) at ~/.julia/packages/ChainRules/3yDBX/src/rulesets/LinearAlgebra/structured.jl:55
...
Stacktrace:
...
```

My intention is to diagonalize a multidimensional array (that depends on some trainable parameters) along a specific axis before operating in further. Any suggestions for debugging/rewriting this?