import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Canberra Focal Loss implementation.
    Computes the Canberra distance with focal loss modulation between two sets of vectors.
    """
    def __init__(self, alpha=1.0, gamma=2.0):
        super(Model, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
    
    def forward(self, x: torch.Tensor, y: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """
        Compute Canberra Focal Loss 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]
            target (torch.Tensor): Target weights [batch_size]

        Returns:
            torch.Tensor: Canberra Focal Loss [batch_size]
        """
        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 element-wise Canberra distance terms
        diff = torch.abs(x - y)
        denom = torch.abs(x) + torch.abs(y)
        canberra_terms = torch.where(denom > 0, diff / denom, torch.zeros_like(diff))
        
        # Apply focal modulation element-wise using built-in functions
        focal_weights = self.alpha * torch.pow(1.0 + canberra_terms, -self.gamma)
        focal_terms = focal_weights * canberra_terms
        
        # Sum along feature dimension using built-in function
        focal_loss = torch.sum(focal_terms, dim=1)
        
        return focal_loss

batch_size = 256
feature_dim = 1024

def get_inputs():
    x = torch.abs(torch.randn(batch_size, feature_dim)) + 0.1
    y = torch.abs(torch.randn(batch_size, feature_dim)) + 0.1
    target = torch.ones(batch_size)
    return [x, y, target]

def get_init_inputs():
    return [1.0, 2.0]
