Speed Comparison Python v Julia for custom layers

[gdalle]

@willleeney

Currently these two implementations clock in at – Python: 146 μs – Julia: 34.363 μs

I think it would be fairer to measure without the random number generation, although it is probably negligible.

@willleeney

Is it fairer to compare backward pass as well? I am primarily interested in speed for inference and secondly memory overhead, which is why I didn’t include any optimisation

In this use case I think you’re okay with just the forward pass, at least if you work with a fully trained network from the start.

@willleeney

function (m::CustomLayer)(xi, xj)
    return transpose(m.σ(((xi * m.Wq) * transpose(xj * m.Wk)) ./ m.dk, dims=2) * (xj * m.Wv))
end

If you want to get peak performance from your Julia code, you may want to think about allocations. Maybe there is a way to reuse some storage in this computation?
Of course this will also mean that the backward pass with Zygote will fail due to mutation, but if you don’t need it who cares