import torch
from safetensors.torch import load_file

def load_model(path='model.safetensors'):
    return load_file(path)

def multiply_3x3_ref(a, b):
    """Reference 3x3 multiplier."""
    return a * b

if __name__ == '__main__':
    print('Wallace Tree 3x3 Multiplier')
    print('A[2:0] x B[2:0] = P[5:0]')
    print()
    print('Examples:')
    for a in range(8):
        for b in [1, 7]:
            p = a * b
            print(f'  {a} x {b} = {p}')