import torch
from torch import device
import triton
import triton.language as tl
from torch._inductor.runtime.triton_heuristics import grid
from torch._C import _cuda_getCurrentRawStream as get_raw_stream
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor


@triton.jit
def triton_poi_fused_constant_pad_nd_0(in_ptr0, out_ptr0, xnumel, XBLOCK:
    tl.constexpr):
    xnumel = 1024
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 8 % 8
    x0 = xindex % 8
    x2 = xindex // 64
    x3 = xindex
    tmp0 = x1
    tmp1 = tl.full([1], 4, tl.int64)
    tmp2 = tmp0 < tmp1
    tmp3 = x0
    tmp4 = tmp3 < tmp1
    tmp5 = tmp2 & tmp4
    tmp6 = tl.load(in_ptr0 + (x0 + 4 * x1 + 16 * x2), tmp5 & xmask, other=0.0)
    tl.store(out_ptr0 + x3, tmp6, xmask)


@triton.jit
def triton_poi_fused_add_1(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.
    constexpr):
    xnumel = 512
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 8
    x1 = xindex // 8 % 4
    x4 = xindex // 32
    x5 = xindex % 128
    x6 = xindex
    tmp0 = tl.load(in_ptr0 + (x0 + 10 * x1 + 80 * x4), xmask)
    tmp1 = tl.load(in_ptr1 + x5, xmask, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tl.store(out_ptr0 + x6, tmp2, xmask)


def call(args):
    primals_1, primals_2, primals_3 = args
    args.clear()
    assert_size_stride(primals_1, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(primals_2, (4, 4, 4, 4, 2), (128, 32, 8, 2, 1))
    assert_size_stride(primals_3, (4, 4, 4, 2), (32, 8, 2, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4, 8, 8), (256, 64, 8, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_constant_pad_nd_0[grid(1024)](primals_1, buf0, 
            1024, XBLOCK=128, num_warps=4, num_stages=1)
        del primals_1
        buf1 = torch.ops.aten._fft_r2c.default(buf0, [2, 3], 0, True)
        del buf0
        buf2 = buf1
        del buf1
        buf3 = torch.ops.aten.full.default([4, 4, 8, 5], 0, dtype=torch.
            complex64, layout=torch.strided, device=device(type='cuda',
            index=0), pin_memory=False)
        buf4 = buf3
        del buf3
        buf5 = torch.ops.aten.slice.Tensor(buf2, 2, 0, 4)
        buf6 = buf5
        buf7 = torch.ops.aten.slice.Tensor(buf6, 3, 0, 4)
        buf8 = buf7
        buf9 = torch.ops.aten.view_as_complex.default(primals_2)
        buf10 = buf9
        buf11 = torch.ops.aten.unsqueeze.default(buf8, 4)
        buf12 = buf11
        buf13 = torch.ops.aten.permute.default(buf12, [0, 4, 2, 3, 1])
        buf14 = buf13
        buf15 = torch.ops.aten.unsqueeze.default(buf10, 4)
        buf16 = buf15
        buf17 = torch.ops.aten.permute.default(buf16, [4, 1, 2, 3, 0])
        buf18 = buf17
        buf19 = torch.ops.aten.permute.default(buf14, [2, 3, 0, 4, 1])
        buf20 = buf19
        buf21 = torch.ops.aten.clone.default(buf20, memory_format=torch.
            contiguous_format)
        del buf11
        del buf12
        del buf13
        del buf14
        del buf19
        del buf2
        del buf20
        del buf5
        del buf6
        del buf7
        del buf8
        buf22 = buf21
        del buf21
        buf23 = torch.ops.aten.reshape.default(buf22, [16, 4, 4])
        buf24 = buf23
        buf25 = torch.ops.aten.permute.default(buf18, [2, 3, 4, 1, 0])
        buf26 = buf25
        buf27 = torch.ops.aten.reshape.default(buf26, [16, 4, 4])
        buf28 = buf27
        buf29 = torch.ops.aten.bmm.default(buf24, buf28)
        del buf10
        del buf15
        del buf16
        del buf17
        del buf18
        del buf25
        del buf26
        del buf27
        del buf28
        del buf9
        del primals_2
        buf30 = buf29
        del buf29
        buf31 = torch.ops.aten.reshape.default(buf30, [4, 4, 4, 1, 4])
        buf32 = buf31
        buf33 = torch.ops.aten.permute.default(buf32, [2, 4, 0, 1, 3])
        buf34 = buf33
        buf35 = torch.ops.aten.reshape.default(buf34, [4, 4, 4, 4])
        buf36 = buf35
        buf37 = torch.ops.aten.slice.Tensor(buf4, 2, 0, 4)
        buf38 = buf37
        buf39 = torch.ops.aten.slice.Tensor(buf38, 3, 0, 4)
        buf40 = buf39
        buf41 = torch.ops.aten.copy.default(buf40, buf36)
        del buf30
        del buf31
        del buf32
        del buf33
        del buf34
        del buf35
        del buf36
        buf42 = buf41
        del buf41
        buf43 = torch.ops.aten.slice.Tensor(buf4, 2, 0, 4)
        buf44 = buf43
        buf45 = torch.ops.aten.slice_scatter.default(buf44, buf42, 3, 0, 4)
        del buf42
        buf46 = buf45
        del buf45
        buf47 = torch.ops.aten.slice_scatter.default(buf4, buf46, 2, 0, 4)
        del buf46
        buf48 = buf47
        del buf47
        buf49 = torch.ops.aten.view_as_complex.default(primals_3)
        buf50 = buf49
        buf51 = torch.ops.aten.slice.Tensor(buf48, 2, 0, 4)
        buf52 = buf51
        buf53 = torch.ops.aten.slice.Tensor(buf52, 3, 0, 4)
        buf54 = buf53
        buf55 = torch.ops.aten.view.dtype(buf54, torch.float32)
        buf56 = buf55
        buf57 = torch.ops.aten.view.dtype(buf50, torch.float32)
        buf58 = buf57
        buf59 = empty_strided_cuda((4, 4, 4, 4, 2), (128, 32, 8, 2, 1),
            torch.float32)
        triton_poi_fused_add_1[grid(512)](buf56, buf58, buf59, 512, XBLOCK=
            256, num_warps=4, num_stages=1)
        del buf49
        del buf50
        del buf51
        del buf52
        del buf53
        del buf54
        del buf55
        del buf56
        del buf57
        del buf58
        del primals_3
        buf60 = torch.ops.aten.view.dtype(reinterpret_tensor(buf59, (4, 4, 
            4, 8), (128, 32, 8, 1), 0), torch.complex64)
        buf61 = buf60
        buf62 = torch.ops.aten.slice.Tensor(buf48, 2, 0, 4)
        buf63 = buf62
        buf64 = torch.ops.aten.slice_scatter.default(buf63, buf61, 3, 0, 4)
        del buf59
        del buf60
        del buf61
        del buf62
        del buf63
        buf65 = buf64
        del buf64
        buf66 = torch.ops.aten.slice_scatter.default(buf48, buf65, 2, 0, 4)
        del buf48
        del buf65
        buf67 = buf66
        del buf66
        buf68 = torch.ops.aten.slice.Tensor(buf67, 2, 0, 4)
        buf69 = buf68
        buf70 = torch.ops.aten.slice.Tensor(buf69, 3, 0, 4)
        buf71 = buf70
        buf72 = torch.ops.aten.slice.Tensor(buf67, 2, 0, 4)
        buf73 = buf72
        buf74 = torch.ops.aten.slice_scatter.default(buf73, buf71, 3, 0, 4)
        del buf68
        del buf69
        del buf70
        del buf71
        del buf72
        del buf73
        buf75 = buf74
        del buf74
        buf76 = torch.ops.aten.slice_scatter.default(buf67, buf75, 2, 0, 4)
        del buf67
        del buf75
        buf77 = buf76
        del buf76
        buf78 = torch.ops.aten._fft_c2r.default(buf77, [2, 3], 2, 8)
        del buf77
        buf79 = buf78
        del buf78
        buf80 = torch.ops.aten.permute.default(buf24, [0, 2, 1])
        buf81 = buf80
        buf82 = torch.ops.aten._conj.default(buf81)
        buf83 = buf82
    return reinterpret_tensor(buf79, (4, 4, 4, 4), (256, 64, 8, 1), 0
        ), buf4, buf83


class FourierConv2dNew(torch.nn.Module):

    def __init__(self, in_channels, out_channels, size_x, size_y, bias=True,
        periodic=False):
        super(FourierConv2dNew, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        if not periodic:
            self.size_x = size_x
            self.size_y = size_y
        else:
            self.size_x = size_x // 2
            self.size_y = size_y // 2
        self.weights = torch.nn.Parameter(torch.view_as_real(1 / (
            in_channels * out_channels) * torch.rand(in_channels,
            out_channels, self.size_x, self.size_y, dtype=torch.cfloat)))
        self.biases = torch.nn.Parameter(torch.view_as_real(1 /
            out_channels * torch.rand(out_channels, self.size_x, self.
            size_y, dtype=torch.cfloat)))
        self.bias = bias
        self.periodic = periodic

    def forward(self, input_0):
        primals_2 = self.weights
        primals_3 = self.biases
        primals_1 = input_0
        output = call([primals_1, primals_2, primals_3])
        return output[0]


Measured latency_ms: 0.190454