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


class Model(nn.Module):
    def __init__(self, target):
        super(Model, self).__init__()
        self.target = nn.Parameter(target)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        mean_x = x.mean(dim=-1, keepdim=True)
        mean_t = self.target.mean(dim=-1, keepdim=True)

        xm = x - mean_x
        tm = self.target - mean_t

        cov = torch.sum(xm * tm, dim=-1)
        sx = torch.sqrt(torch.sum(xm * xm, dim=-1))
        st = torch.sqrt(torch.sum(tm * tm, dim=-1))

        correlation = cov / (sx * st + 1e-6)

        return F.elu(correlation)


batch_size = 128
input_dim = 1024


def get_inputs():
    x = torch.randn(batch_size, input_dim)
    return [x]


def get_init_inputs():
    target = torch.randn(input_dim)
    return [target]