import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Computes the pairwise Euclidean distance matrix between two sets of vectors.
    This implementation uses broadcasting, which is memory-intensive.
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x (torch.Tensor): A tensor of shape (N, D).
            y (torch.Tensor): A tensor of shape (M, D).

        Returns:
            torch.Tensor: A tensor of shape (N, M) where Z[i, j] is the distance
                          between x[i] and y[j].
        """
        # x.unsqueeze(1) -> (N, 1, D)
        # y.unsqueeze(0) -> (1, M, D)
        # The result of subtraction is (N, M, D)
        diff = x.unsqueeze(1) - y.unsqueeze(0)
        # Compute L2 norm along the last dimension
        dist_matrix = torch.norm(diff, p=2, dim=-1)
        return dist_matrix

# 测试数据配置
N = 1024  # Number of vectors in the first set
M = 1024  # Number of vectors in the second set
D = 512   # Dimensionality of each vector

def get_inputs():
    x = torch.randn(N, D)
    y = torch.randn(M, D)
    return [x, y]

def get_init_inputs():
    return [] # No special initialization inputs needed
