<jupyter_start><jupyter_text><jupyter_code>%matplotlib inline
from torch import tensor
import numpy as np
import torch
def test_roughly_eq(m1, m2, tolerance=1e-5):
  assert ((m1-m2).abs() > tolerance).sum() == 0, f"Expected:\n {m1}\n But actual was: \n {m2}"


def check_can_multiply_3x1_and_1x3(func):
  tensor1 = torch.arange(0, 3, 1.).unsqueeze(1)
  tensor2 = (torch.arange(0, 3, 1.)*10).unsqueeze(0)
  torch_mul = tensor1 @ tensor2
  own_mul = func(tensor1, tensor2)

  test_roughly_eq(torch_mul, own_mul)

def check_can_multiply_3x2_and_2x3(func):
  tensor1 = torch.arange(0, 3, 1.).unsqueeze(1).repeat((1, 2))
  tensor2 = (torch.arange(0, 3, 1.)*10).unsqueeze(0).repeat((2, 1))
  torch_mul = tensor1 @ tensor2
  own_mul = func(tensor1, tensor2)

  test_roughly_eq(torch_mul, own_mul)
  
def run_tests(func):
  check_can_multiply_3x1_and_1x3(func)
  check_can_multiply_3x2_and_2x3(func)

def matmul_1(ltensor, rtensor):
  result = torch.zeros(ltensor.shape[0], rtensor.shape[1])
  for lrow in range(ltensor.shape[0]):
    for rcol in range(rtensor.shape[1]):
      for k in range(ltensor.shape[1]):
        result[lrow, rcol] += ltensor[lrow, k] * rtensor[k, rcol];
  return result;

run_tests(matmul_1)
def matmul_2(ltensor, rtensor):
  result = torch.zeros(ltensor.shape[0], rtensor.shape[1])
  for lrow in range(ltensor.shape[0]):
    for rcol in range(rtensor.shape[1]):
      result[lrow, rcol] = (ltensor[lrow, :] * rtensor[:, rcol]).sum();
  return result;

run_tests(matmul_2)
def matmul_3(ltensor, rtensor):
  result = torch.zeros(ltensor.shape[0], rtensor.shape[1])
  for lrow in range(ltensor.shape[0]):
    # result[lrow, rcol] = (ltensor[lrow, :] * rtensor[:, rcol]).sum();
    result[lrow, :] = (ltensor[lrow, :].unsqueeze(-1) * rtensor).sum(dim=0)
  return result;

run_tests(matmul_3)
def matmul_4(ltensor, rtensor):
  return torch.einsum('ij,jk -> ik', ltensor, rtensor);

run_tests(matmul_4)
import timeit
from collections import defaultdict
series = defaultdict(list)

tested_functions = {
    'for loops': matmul_1,
    'broadcast1': matmul_2,
    'broadcast2': matmul_3,
    'einsum': matmul_4,
    'torch': lambda m,n: m @ n
}

for size in [10, 100, 150, 200]:
  print(f'\n= Testing size {size}')
  series['size'].append(size)  
  tensor1 = torch.arange(0, size, 1.).unsqueeze(1)
  tensor2 = (torch.arange(0, size, 1.)*10).unsqueeze(0)
  for name, f in tested_functions.items():
    print(f'Timing {name}')
    elapsed = timeit.timeit('f(tensor1, tensor2)', globals=globals(), number=15)
    series[name].append(elapsed)
import pandas as pd
sizes = series['size']
del series['size']
df = pd.DataFrame(series, index=sizes)
df
df.plot(figsize=(14, 7), logy=True)<jupyter_output><empty_output><jupyter_text>Named tensor<jupyter_code>!pip install git+https://github.com/harvardnlp/namedtensor -q
from namedtensor import ntorch
a = ntorch.tensor(torch.arange(0, 3, 1.).unsqueeze(1).repeat((1, 2)), names=('rows', 'cols'))
b = ntorch.tensor((torch.arange(0, 3, 1.)*10).unsqueeze(0).repeat((2, 1)), names=('rows', 'cols'))
b
a[{'rows': 0}]
b[{'rows': 0}]
# Note to self: elementwise access looks ugly and not really designed for
a[{'rows': 0, 'cols': 0}].numpy() * b[{'rows': 0, 'cols': 0}].numpy()
a[{'rows': 0, 'cols': 0}].numpy()
def named_mul_1(ltensor, rtensor):
  result = torch.zeros(ltensor.shape['rows'], rtensor.shape['cols'])
  for lrow in range(ltensor.shape['rows']):
    for rcol in range(rtensor.shape['cols']):
      for k in range(ltensor.shape['cols']):
        result[lrow, rcol] += ltensor[{'rows': lrow, 'cols': k}].numpy() * rtensor[{'rows': k, 'cols': rcol}].numpy();
  return result;

named_mul_1(a, b)
a.shape
a
b
# Note to self: broadcasting looks nicer in writing, but in here both dims are broadcasted which is a bummer, especially given unsqueze ban
a[{'rows':1}] * b[{'cols': 2}]
a[{'rows':1}].numpy() * b[{'cols': 2}].numpy()

(a[{'rows':1}] * b[{'cols': 2}])
a[{'rows':1}]
b[{'cols': 2}]
a[{'rows':1}] * b[{'cols': 2}]
def named_mul_2(ltensor, rtensor):
  result = ntorch.zeros((ltensor.shape['rows'], rtensor.shape['cols']), names=('cols', 'rows'))
  any_ = 0
  for lrow in range(ltensor.shape['rows']):
    for rcol in range(rtensor.shape['cols']):
      # Note to self: scalars are so unintuitive Dimension () does not exist. Available dimensions are ()
      broadcasted_on_both_dims = (ltensor[{'rows': lrow}] * rtensor[{'cols': rcol}])
      single_dim_result = broadcasted_on_both_dims[{'rows': any_}]
      result[{'rows': lrow, 'cols': rcol}] = int(single_dim_result.sum().numpy())
  return result;

named_mul_2(a, b)
# result[lrow, rcol] = (ltensor[lrow, :] * rtensor[:, rcol]).sum();
# result[lrow, :] = (ltensor[lrow, :].unsqueeze(-1) * rtensor).sum(dim=0)
# Note to self: I guess that's the end of it, this sort of encapsulation makes it difficult for two different dimensions to interact looks good on images though
a[{'rows': 1}] * b<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.18255578093306288, 'fraction_numerical': 0.03184584178498986, 'mean_word_length': 5.133084577114428, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 11, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '518959', 'license': 'Apache-2.0', 'n_tokens_mistral': 2000, 'n_tokens_neox': 1948, 'n_words': 433, 'path': 'fastai3_part2_00_matmul.ipynb', 'repo_name': 'cluePrints/course-v3', 'size': '4930'}