import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Minkowski Distance implementation.
    Computes the Minkowski distance between two sets of vectors with order p.
    """
    def __init__(self, p=2):
        super(Model, self).__init__()
        self.p = p
        if p <= 0:
            raise ValueError("p must be positive")
    
    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        """
        Compute Minkowski distance between x and y.

        Args:
            x (torch.Tensor): First set of vectors [batch_size, feature_dim]
            y (torch.Tensor): Second set of vectors [batch_size, feature_dim]

        Returns:
            torch.Tensor: Minkowski distances [batch_size]
        """
        # Input validation
        if x.shape != y.shape:
            raise ValueError(f"Input tensors must have the same shape, got {x.shape} and {y.shape}")
        
        if x.dim() != 2:
            raise ValueError(f"Input tensors must be 2D, got {x.dim()}D")
        
        # Compute absolute differences
        abs_diff = torch.abs(x - y)
        
        # Compute Minkowski distance: (Σ|x_i - y_i|^p)^(1/p)
        if self.p == 1:
            # Manhattan distance
            minkowski_dist = torch.sum(abs_diff, dim=1)
        elif self.p == 2:
            # Euclidean distance
            minkowski_dist = torch.sqrt(torch.sum(abs_diff ** 2, dim=1))
        elif self.p == float('inf'):
            # Chebyshev distance
            minkowski_dist = torch.max(abs_diff, dim=1)[0]
        else:
            # General Minkowski distance
            minkowski_dist = torch.pow(torch.sum(torch.pow(abs_diff, self.p), dim=1), 1.0/self.p)
        
        return minkowski_dist

batch_size = 256
feature_dim = 512

def get_inputs():
    # Generate two sets of vectors
    x = torch.randn(batch_size, feature_dim)
    y = torch.randn(batch_size, feature_dim)
    return [x, y]

def get_init_inputs():
    return [2]  # p value (default: Euclidean distance)
