import torch
import torch.nn as nn

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.scale = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x):
        var = x.pow(2).mean(-1, keepdim=True)
        norm_x = x * torch.rsqrt(var + self.eps)
        return norm_x * self.scale

class Model(nn.Module):
    def __init__(self, dim):
        super(Model, self).__init__()
        self.rmsnorm = RMSNorm(dim)

    def forward(self, x, y):
        eps = 1e-6
        val_x = torch.abs(x) + eps
        val_y = torch.abs(y) + eps
        dist = torch.abs(torch.log(val_x) - torch.log(val_y))
        out = self.rmsnorm(dist)
        return out.mean()

batch_size = 16
input_dim = 1024

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

def get_init_inputs():
    return [input_dim]