This is a MWE for a problem i’m having, I’ve adapted matrix multiplication to accept `trans` arguments telling whether the matrix is to be transposed.
As for the adjoint definition, I looked at `NNlib.batched_mul`

When I try to get the gradient of x.x^T w.r.t. x using `my_mul`, get an error shown below.

``````x = rand(2, 3)

function my_mul(A, B, transA, transB) #transA/transB is either 'N' or 'T'
(transA == 'N' ? A : A')*(transB == 'N' ? B : B')
end

my_mul(x, x, 'N', 'T')

using Zygote

Zygote.@adjoint function my_mul(A, B, transA, transB)
d = Dict('N'=>'T', 'T'=>'N')

C = my_mul(A, B, transA, transB)
C, Δ -> (my_mul(Δ, B, 'N', d[transB]), my_mul(A, Δ, d[transA], 'N'))
end

sum(x*x')
end # expected result

sum(my_mul(x, x, 'N', 'T'))
end # unexpected!
``````
``````ERROR: DimensionMismatch("arrays could not be broadcast to a common size; got a dimension with lengths 2 and 3")
Stacktrace:
[7] accum(::Array{Float64,2}, ::Array{Float64,2})
at C:\Users\jules\.julia\packages\Zygote\Xgcgs\src\lib\lib.jl:16
[8] (::Zygote.var"#41#42"{typeof(∂(#13))})(::Float64) at C:\Users\jules\.julia\packages\Zygote\Xgcgs\src\compiler\interface.jl:45
[10] top-level scope at REPL[1]:1
I believe the failed broadcasting is coming from `Zygote.accum` which tries to add both the adjoint results together.
Note that this only happens if I’m multiplying an array with itself, that is, replacing one of the `x`'s with a copy fixes it.
When I take the gradient of `sum(x*x')` I do get a sensible answer.