from __future__ import annotations

import torch


torch.manual_seed(0)


def verify(x1_shape: tuple[int, ...], x2_shape: tuple[int, ...]) -> None:
    assert x1_shape[-1] == x2_shape[-1]
    x1 = torch.randn(x1_shape, dtype=torch.float64)
    x2 = torch.randn(x2_shape, dtype=torch.float64)
    w = torch.rand(x1_shape[-1], dtype=torch.float64)
    sqd = (x1.unsqueeze(-2) - x2.unsqueeze(-3)).square_()
    ans = (w * sqd).sum(dim=-1)
    out = sqd.matmul(w)
    assert torch.allclose(ans, out)
    print(f"{x1_shape=}, {x2_shape}: Successfull")


for i in range(30):
    print(f"Seed{i}")
    for x1_shape in [(4,), (3, 4), (6, 3, 4)]:
        for x2_shape in [(7, 4), (6, 7, 4)]:
            verify(x1_shape, x2_shape)