import torch
import torch.nn as nn
import torch.nn.functional as F

# 定义维度常量
N, C, H, W = 32, 64, 56, 56
EPS = 1e-6


class Tanimoto(nn.Module):


    def __init__(self, eps=1e-6):
        super().__init__()
        self.eps = eps
        # 我们将在 C, H, W 维度上进行归约
        self.reduction_dims = (1, 2, 3)

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:

        x_dot_y = torch.sum(x * y, dim=self.reduction_dims)


        x_norm_sq = torch.sum(x * x, dim=self.reduction_dims)
        y_norm_sq = torch.sum(y * y, dim=self.reduction_dims)


        denominator = x_norm_sq + y_norm_sq - x_dot_y


        similarity = (x_dot_y + self.eps) / (denominator + self.eps)

        return similarity


class Model(nn.Module):


    def __init__(self):
        super().__init__()
        self.op = Tanimoto(EPS)

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        return self.op(x, y)


def get_inputs():

    x = torch.randn(N, C, H, W, dtype=torch.float32)
    y = torch.randn(N, C, H, W, dtype=torch.float32)
    return [x, y]


def get_init_inputs():

    return []
