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


BATCH = 64
CHANNELS = 64
WIDTH_IN = 4096
WIDTH_OUT = 4096

class Model(nn.Module):
    def __init__(self):
        super().__init__()

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

        x_4d = x.unsqueeze(2) 
        
      
        grid_x = grid.unsqueeze(1)
  
        grid_y = torch.zeros_like(grid_x)
        
    
        grid_2d = torch.cat([grid_x, grid_y], dim=-1)   

        out_4d = F.grid_sample(x_4d, grid_2d, align_corners=True, mode='bilinear', padding_mode='zeros')

        return out_4d.squeeze(2)

def get_inputs():

    w = torch.arange(WIDTH_IN, device='cuda', dtype=torch.float32).view(1, 1, WIDTH_IN)
    x = w.expand(BATCH, CHANNELS, WIDTH_IN).contiguous()

    grid_x = torch.linspace(-1, 1, WIDTH_OUT, device='cuda', dtype=torch.float64).float()

    grid = grid_x.view(1, WIDTH_OUT, 1).expand(BATCH, -1, -1).contiguous()
    
    return [x, grid]

def get_init_inputs():
    return []