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 JaccardSimilarity(nn.Module):

    def __init__(self, eps=1e-6):
        super().__init__()
        self.eps = eps

    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        intersection = torch.sum(x1 * x2, dim=[1, 2, 3])

        sum_sq1 = torch.sum(x1 * x1, dim=[1, 2, 3])
        sum_sq2 = torch.sum(x2 * x2, dim=[1, 2, 3])

        union = sum_sq1 + sum_sq2 - intersection

        return intersection / (union + self.eps)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.op = JaccardSimilarity(EPS)

    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        return self.op(x1, x2)


def get_inputs():
    x1 = torch.randn(N, C, H, W, dtype=torch.float32)
    x2 = torch.randn(N, C, H, W, dtype=torch.float32)
    return [x1, x2]


def get_init_inputs():
    return []