import torch
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
from torch._inductor.runtime.triton_helpers import math as tl_math
import torch.nn as nn
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_cat_cos_mul_sin_0(in_ptr0, out_ptr0, out_ptr1,
    out_ptr2, out_ptr3, out_ptr4, out_ptr5, out_ptr6, out_ptr7, out_ptr8,
    xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x0 = xindex % 4
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp1 = 1.0
    tmp2 = tmp0 * tmp1
    tmp3 = tl_math.sin(tmp2)
    tmp4 = tl_math.cos(tmp2)
    tmp5 = 2.5198421478271484
    tmp6 = tmp0 * tmp5
    tmp7 = tl_math.sin(tmp6)
    tmp8 = tl_math.cos(tmp6)
    tmp9 = 6.349603652954102
    tmp10 = tmp0 * tmp9
    tmp11 = tl_math.sin(tmp10)
    tmp12 = tl_math.cos(tmp10)
    tmp13 = 16.0
    tmp14 = tmp0 * tmp13
    tmp15 = tl_math.sin(tmp14)
    tmp16 = tl_math.cos(tmp14)
    tl.store(out_ptr0 + (x0 + 36 * x1), tmp0, xmask)
    tl.store(out_ptr1 + (x0 + 36 * x1), tmp3, xmask)
    tl.store(out_ptr2 + (x0 + 36 * x1), tmp4, xmask)
    tl.store(out_ptr3 + (x0 + 36 * x1), tmp7, xmask)
    tl.store(out_ptr4 + (x0 + 36 * x1), tmp8, xmask)
    tl.store(out_ptr5 + (x0 + 36 * x1), tmp11, xmask)
    tl.store(out_ptr6 + (x0 + 36 * x1), tmp12, xmask)
    tl.store(out_ptr7 + (x0 + 36 * x1), tmp15, xmask)
    tl.store(out_ptr8 + (x0 + 36 * x1), tmp16, xmask)


def call(args):
    arg0_1, = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 4, 4), (64, 16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf9 = empty_strided_cuda((4, 4, 4, 36), (576, 144, 36, 1), torch.
            float32)
        buf0 = reinterpret_tensor(buf9, (4, 4, 4, 4), (576, 144, 36, 1), 0)
        buf1 = reinterpret_tensor(buf9, (4, 4, 4, 4), (576, 144, 36, 1), 4)
        buf2 = reinterpret_tensor(buf9, (4, 4, 4, 4), (576, 144, 36, 1), 8)
        buf3 = reinterpret_tensor(buf9, (4, 4, 4, 4), (576, 144, 36, 1), 12)
        buf4 = reinterpret_tensor(buf9, (4, 4, 4, 4), (576, 144, 36, 1), 16)
        buf5 = reinterpret_tensor(buf9, (4, 4, 4, 4), (576, 144, 36, 1), 20)
        buf6 = reinterpret_tensor(buf9, (4, 4, 4, 4), (576, 144, 36, 1), 24)
        buf7 = reinterpret_tensor(buf9, (4, 4, 4, 4), (576, 144, 36, 1), 28)
        buf8 = reinterpret_tensor(buf9, (4, 4, 4, 4), (576, 144, 36, 1), 32)
        get_raw_stream(0)
        triton_poi_fused_cat_cos_mul_sin_0[grid(256)](arg0_1, buf0, buf1,
            buf2, buf3, buf4, buf5, buf6, buf7, buf8, 256, XBLOCK=128,
            num_warps=4, num_stages=1)
        del arg0_1
    return buf9,


class FreqEncoderNew(nn.Module):

    def __init__(self, input_dim, max_freq_log2, N_freqs, log_sampling=True,
        include_input=True, periodic_fns=(torch.sin, torch.cos)):
        super().__init__()
        self.input_dim = input_dim
        self.include_input = include_input
        self.periodic_fns = periodic_fns
        self.output_dim = 0
        if self.include_input:
            self.output_dim += self.input_dim
        self.output_dim += self.input_dim * N_freqs * len(self.periodic_fns)
        if log_sampling:
            self.freq_bands = 2.0 ** torch.linspace(0.0, max_freq_log2, N_freqs
                )
        else:
            self.freq_bands = torch.linspace(2.0 ** 0.0, 2.0 **
                max_freq_log2, N_freqs)
        self.freq_bands = self.freq_bands.numpy().tolist()

    def forward(self, input_0):
        arg0_1 = input_0
        output = call([arg0_1])
        return output[0]


Measured latency_ms: 0.024387