%load_ext autoreload
%autoreload 2

%matplotlib inline

#export
from exp.nb_00 import *
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,'==')

test_eq(TEST,'test')

# To run tests in console:
# ! python run_notebook.py 01_matmul.ipynb

#export
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

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, x_train.shape, y_train, y_train.shape, 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)

mpl.rcParams['image.cmap'] = 'gray'

img = x_train[0]

img.view(28,28).type()

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

weights = torch.randn(784,10)

bias = torch.zeros(10)

def matmul(a,b):
    ar,ac = a.shape # n_rows * n_cols
    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): # or br
                c[i,j] += a[i,k] * b[k,j]
    return c

m1 = x_valid[:5]
m2 = weights

m1.shape,m2.shape

%time t1=matmul(m1, m2)

t1.shape

len(x_train)

a = tensor([10., 6, -4])
b = tensor([2., 8, 7])
a,b

a + b

(a < b).float().mean()

m = tensor([[1., 2, 3], [4,5,6], [7,8,9]]); m

(m*m).sum().sqrt()

def matmul(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):
            # Any trailing ",:" can be removed
            c[i,j] = (a[i,:] * b[:,j]).sum()
    return c

%timeit -n 10 _=matmul(m1, m2)

890.1/5

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

test_near(t1,matmul(m1, m2))

a

a > 0

a + 1

m

2*m

c = tensor([10.,20,30]); c

m

m.shape,c.shape

m + c

c + m

t = c.expand_as(m)

t

m + t

t.storage()

t.stride(), t.shape

c.unsqueeze(0)

c.unsqueeze(1)

m

c.shape, c.unsqueeze(0).shape,c.unsqueeze(1).shape

c.shape, c[None].shape,c[:,None].shape

c[None].shape,c[...,None].shape

c[:,None].expand_as(m)

m + c[:,None]

c[:,None]

def matmul(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,j] = (a[i,:]          * b[:,j]).sum() # previous
        c[i]   = (a[i  ].unsqueeze(-1) * b).sum(dim=0)
    return c

%timeit -n 10 _=matmul(m1, m2)

885000/277

test_near(t1, matmul(m1, m2))

c[None,:]

c[None,:].shape

c[:,None]

c[:,None].shape

c[None,:] * c[:,None]

c[None] > c[:,None]

# c[i,j] += a[i,k] * b[k,j]
# c[i,j] = (a[i,:] * b[:,j]).sum()
def matmul(a,b): return torch.einsum('ik,kj->ij', a, b)

%timeit -n 10 _=matmul(m1, m2)

885000/55

test_near(t1, matmul(m1, m2))

%timeit -n 10 t2 = m1.matmul(m2)

# time comparison vs pure python:
885000/18

t2 = m1@m2

test_near(t1, t2)

m1.shape,m2.shape

!python notebook2script.py 01_matmul.ipynb