import torch
from torch.utils import benchmark

results = []
for b in [10, 10000, 2000000]:
    for n in [10, 100, 10000, 1000000]:
        if b * n >= 1000000000:
            continue

        description = f'[{b}, {n}]'

        x = torch.rand(b, device='mps')
        y = torch.rand(n, device='mps')

        results.append(benchmark.Timer(
            stmt='torch.outer(x,y)',
            globals={'x': x, 'y': y},
            description=description,
        ).blocked_autorange())

        results.append(benchmark.Timer(
            stmt='torch.einsum("i,j->ij",x,y)',
            globals={'x': x, 'y': y},
            description=description,
        ).blocked_autorange())

compare = benchmark.Compare(results)
compare.trim_significant_figures()
compare.colorize()
compare.print()