%load_ext autoreload
%autoreload 2
%matplotlib inline

import operator

def test(a, b, cmp, cname=None):
    if cname is None: cname=cmp.__name__
    assert cmp(a, b), f"{cname}:\n{a}\n{b}"

def test_eq(a,b): test(a, b, operator.eq, '==')

def near(a,b): return torch.allclose(a, b, rtol=1e-3, atol=1e-5)
def test_near(a,b): test(a,b,near)    

from pathlib import Path
from IPython.core.debugger import set_trace
from fastai import datasets
import pickle, gzip, math, torch, matplotlib as mpl
import matplotlib.pyplot as plt
from torch import tensor
mpl.rcParams['image.cmap'] = 'gray'

MNIST_URL='http://deeplearning.net/data/mnist/mnist.pkl'

path = datasets.download_data(MNIST_URL, ext='.gz')
path

with gzip.open(path, 'rb') as f:
    ((x_train, y_train), (x_valid, y_valid), _) = pickle.load(f, encoding='latin-1')

x_train, y_train, x_valid, y_valid = map(tensor, (x_train, y_train, x_valid, y_valid))
n, c = x_train.shape

x_train.shape, x_train

y_train.shape, y_train

y_train.min(), y_train.max()

assert n == y_train.shape[0] == 50000
test_eq(c, 28*28)
test_eq(y_train.min(), 0)
test_eq(y_train.max(), 9)

img = x_train[0]

plt.imshow(img.view(28, 28))

weights = torch.randn(784, 10)
bias = torch.zeros(10)

m1 = x_valid[:10]
m2 = weights

m1.shape, m2.shape

def matmul1(a, b):
    # a row, a column
    ar, ac = a.shape
    # b row, b column
    br, bc = b.shape
    assert ac == br
    c = torch.zeros(ar, bc)
    
    for i in range(ar):
        for j in range(bc):
            for k in range(ac):
                c[i, j] += a[i, k] * b[k, j]
    return c

%timeit t1 = matmul1(m1, m2)

test_near(m1.matmul(m2), matmul1(m1, m2))

def matmul2(a, b):
    ar, ac = a.shape
    br, bc = b.shape
    assert ac == br
    c = torch.zeros(ar, bc)
    
    for i in range(ar):
        for j in range(bc):
                c[i, j] = (a[i, :] * b[:, j]).sum()
    return c

%timeit t2 = matmul2(m1, m2)

test_near(m1.matmul(m2), matmul2(m1, m2))

def matmul3(a, b):
    ar, ac = a.shape
    br, bc = b.shape
    assert ac == br
    c = torch.zeros(ar, bc)
    
    for i in range(ar):
                c[i, :] = (a[i, :] * b[:, :].t()).sum(1)
    return c

%timeit t3 = matmul3(m1, m2)

test_near(m1.matmul(m2), matmul3(m1, m2))

def matmul4(a, b):
    return torch.einsum('ik,kj->ij', a, b)

%timeit t4 = matmul4(m1, m2)

test_near(m1.matmul(m2), matmul4(m1, m2))

%timeit t5 = m1.matmul(m2)

%timeit t6 = m1 @ m2