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
import 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_add_embedding_0(in_ptr0, in_ptr1, in_ptr2, in_ptr3,
    in_ptr4, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 160
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 10
    x0 = xindex % 10
    x2 = xindex
    tmp0 = tl.load(in_ptr0 + (3 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp8 = tl.load(in_ptr0 + (2 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp17 = tl.load(in_ptr0 + (1 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp26 = tl.load(in_ptr0 + 4 * x1, xmask, eviction_policy='evict_last')
    tmp1 = tmp0.to(tl.int64)
    tmp2 = tl.full([XBLOCK], 24, tl.int32)
    tmp3 = tmp1 + tmp2
    tmp4 = tmp1 < 0
    tmp5 = tl.where(tmp4, tmp3, tmp1)
    tl.device_assert((0 <= tmp5) & (tmp5 < 24) | ~xmask,
        'index out of bounds: 0 <= tmp5 < 24')
    tmp7 = tl.load(in_ptr1 + (x0 + 10 * tmp5), xmask)
    tmp9 = tmp8.to(tl.int64)
    tmp10 = tl.full([XBLOCK], 7, tl.int32)
    tmp11 = tmp9 + tmp10
    tmp12 = tmp9 < 0
    tmp13 = tl.where(tmp12, tmp11, tmp9)
    tl.device_assert((0 <= tmp13) & (tmp13 < 7) | ~xmask,
        'index out of bounds: 0 <= tmp13 < 7')
    tmp15 = tl.load(in_ptr2 + (x0 + 10 * tmp13), xmask)
    tmp16 = tmp7 + tmp15
    tmp18 = tmp17.to(tl.int64)
    tmp19 = tl.full([XBLOCK], 32, tl.int32)
    tmp20 = tmp18 + tmp19
    tmp21 = tmp18 < 0
    tmp22 = tl.where(tmp21, tmp20, tmp18)
    tl.device_assert((0 <= tmp22) & (tmp22 < 32) | ~xmask,
        'index out of bounds: 0 <= tmp22 < 32')
    tmp24 = tl.load(in_ptr3 + (x0 + 10 * tmp22), xmask)
    tmp25 = tmp16 + tmp24
    tmp27 = tmp26.to(tl.int64)
    tmp28 = tl.full([XBLOCK], 13, tl.int32)
    tmp29 = tmp27 + tmp28
    tmp30 = tmp27 < 0
    tmp31 = tl.where(tmp30, tmp29, tmp27)
    tl.device_assert((0 <= tmp31) & (tmp31 < 13) | ~xmask,
        'index out of bounds: 0 <= tmp31 < 13')
    tmp33 = tl.load(in_ptr4 + (x0 + 10 * tmp31), xmask)
    tmp34 = tmp25 + tmp33
    tmp35 = 0.0
    tmp36 = tmp34 + tmp35
    tl.store(out_ptr0 + x2, tmp36, xmask)


def call(args):
    arg0_1, arg1_1, arg2_1, arg3_1, arg4_1 = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 4), (16, 4, 1))
    assert_size_stride(arg1_1, (24, 10), (10, 1))
    assert_size_stride(arg2_1, (7, 10), (10, 1))
    assert_size_stride(arg3_1, (32, 10), (10, 1))
    assert_size_stride(arg4_1, (13, 10), (10, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4, 10), (40, 10, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_add_embedding_0[grid(160)](arg0_1, arg1_1, arg2_1,
            arg3_1, arg4_1, buf0, 160, XBLOCK=256, num_warps=4, num_stages=1)
        del arg0_1
        del arg1_1
        del arg2_1
        del arg3_1
        del arg4_1
    return reinterpret_tensor(buf0, (10, 4, 4), (1, 40, 10), 0),


class CAT_FixedEmbedding(nn.Module):

    def __init__(self, c_in, d_model):
        super(CAT_FixedEmbedding, self).__init__()
        w = torch.zeros(c_in, d_model).float()
        w.require_grad = False
        position = torch.arange(0, c_in).float().unsqueeze(1)
        div_term = (torch.arange(0, d_model, 2).float() * -(math.log(
            10000.0) / d_model)).exp()
        w[:, 0::2] = torch.sin(position * div_term)
        w[:, 1::2] = torch.cos(position * div_term)
        self.emb = nn.Embedding(c_in, d_model)
        self.emb.weight = nn.Parameter(w, requires_grad=False)

    def forward(self, x):
        return self.emb(x).detach()


class CAT_TemporalEmbeddingNew(nn.Module):

    def __init__(self, d_feature=10, embed_type='fixed', freq='h'):
        super(CAT_TemporalEmbeddingNew, self).__init__()
        minute_size = 4
        hour_size = 24
        weekday_size = 7
        day_size = 32
        month_size = 13
        Embed = CAT_FixedEmbedding if embed_type == 'fixed' else nn.Embedding
        if freq == 't':
            self.minute_embed = Embed(minute_size, d_feature)
        self.hour_embed = Embed(hour_size, d_feature)
        self.weekday_embed = Embed(weekday_size, d_feature)
        self.day_embed = Embed(day_size, d_feature)
        self.month_embed = Embed(month_size, d_feature)

    def forward(self, input_0):
        arg1_1 = self.hour_embed.emb.weight
        arg2_1 = self.weekday_embed.emb.weight
        arg3_1 = self.day_embed.emb.weight
        arg4_1 = self.month_embed.emb.weight
        arg0_1 = input_0
        output = call([arg0_1, arg1_1, arg2_1, arg3_1, arg4_1])
        return output[0]


Measured latency_ms: 0.052703