import torch
import torch.nn as nn

torch.backends.cuda.matmul.allow_tf32 = False

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.H = 512
        self.W = 512
        self.B = 4

    def forward(self, voxel_features: torch.Tensor, coords: torch.Tensor) -> torch.Tensor:
        C = voxel_features.shape[1]
        canvas = torch.zeros(self.B, C, self.H, self.W, dtype=voxel_features.dtype, device=voxel_features.device)
        
        batch_idx = coords[:, 0].long()
        y_idx = coords[:, 2].long()
        x_idx = coords[:, 3].long()
        
        canvas[batch_idx, :, y_idx, x_idx] = voxel_features
        return canvas

N = 64000
C = 64
H = 512
W = 512
B = 4

def get_inputs():
    feats = torch.randn(N, C, device='cuda', dtype=torch.float32)
    
    total_voxels = B * H * W
    indices = torch.randperm(total_voxels, device='cuda')[:N]
    
    b_idx = indices // (H * W)
    rem = indices % (H * W)
    y_idx = rem // W
    x_idx = rem % W
    z_idx = torch.zeros_like(b_idx)
    
    coords = torch.stack([b_idx, z_idx, y_idx, x_idx], dim=1).int()
    return [feats, coords]

def get_init_inputs():
    return []