speed test

[Xylambda]

import torch
import perfplot
import numpy as np
import toydiff as tdf


def matmul_torch(t1, t2):
    return torch.matmul(t1, t2)


def matmul_avagrad(t1, t2):
    return tdf.matmul(t1, t2)


def data(n):
    arr = np.random.rand(n, n)
    tensor_t = torch.from_numpy(arr)
    tensor_a = tdf.Tensor(arr)
    
    return tensor_t, tensor_a


perfplot.show(
    setup=lambda n: data(n),
    kernels=[
        lambda tt, ta: matmul_torch(tt, tt),
        lambda tt, ta: matmul_avagrad(ta, ta),
    ],
    labels=["PyTorch", "Avagrad"],
    n_range=[2**k for k in range(14)],
    xlabel="len(a)",
)