New to Julia and looking at comparisons of speed in forward pass with pytorch. I want to optimise performance of training a neural network that will contain custom layers. I have provide some code for an example. I want to know how fast I can speed up my Julia code. Currently these two implementations clock in at – Python: 146 μs – Julia: 34.363 μs – which suggests that Julia is much faster. Please can someone help me to fairly compare the two languages for speed. 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 of the network in the examples. Also would these results change as the layers become sparser or as the size of the layers scale?

```
using Flux
using Statistics
using LinearAlgebra
using GraphNeuralNetworks
struct CustomLayer <: GNNLayer
Wq::AbstractMatrix
Wk::AbstractMatrix
Wv::AbstractMatrix
dk::Float64
σ::typeof(softmax)
end
Flux.@functor CustomLayer
function CustomLayer(n_in::Int, n_out::Int, n_neighbors::Int)
Wq = Flux.glorot_uniform(n_in, n_neighbors)
Wk = Flux.glorot_uniform(n_in, n_neighbors)
Wv = Flux.glorot_uniform(n_in, n_out)
dk = √n_in
σ = softmax
CustomLayer(Wq, Wk, Wv, dk, σ)
end
function (m::CustomLayer)(xi, xj)
return transpose(m.σ(((xi * m.Wq) * transpose(xj * m.Wk)) ./ m.dk, dims=2) * (xj * m.Wv))
end
myLayer = CustomLayer(1000, 10, 20)
function test_function()
xi = rand(Float32, 1, 1000)
xj = rand(Float32, 10, 1000)
out = myLayer(xi, xj)
return
end
@btime test_function()
```

```
import torch
import torch.nn as nn
import timeit
import numpy as np
class CustomLayer(nn.Module):
def __init__(self, n_in, n_out, n_neighbors):
super().__init__()
self.Wq = nn.Linear(n_in, n_neighbors)
self.Wk = nn.Linear(n_in, n_neighbors)
self.Wv = nn.Linear(n_in, n_out)
self.dk = np.sqrt(n_in)
self.softmax = nn.Softmax(dim=-1)
def forward(self, xi, xj):
attention = self.softmax(self.Wq(xi) @ self.Wk(xj).T / self.dk)
return attention @ self.Wv(xj)
myLayer = CustomLayer(1000, 10, 20)
def test_function():
xi = torch.rand(1, 1000)
xj = torch.rand(10, 1000)
out = myLayer(xi, xj)
return
number_exc = 10000
print(timeit.timeit("test_function()", setup="from __main__ import test_function", number=number_exc) / number_exc)
```