AutoGrad benchmark of matmul

autograd_benchmark.py

import timeit
import autograd.numpy as np
from autograd import grad


def fn(a, b):
    return np.sum(a @ b)


n = 1000
a = np.random.rand(n, n)
b = np.random.rand(n, n)

grad_fn = grad(fn)
t = min(timeit.repeat('grad_fn(a, b)', globals=globals(), number=1, repeat=10))
print("grad_fn(a, b) is %d ms" % (t * 1000))

t = min(timeit.repeat('fn(a, b)', globals=globals(), number=1, repeat=10))
print("fn(a, b) is %d ms" % (t * 1000))