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
from torch._inductor.runtime import triton_helpers
from torch._inductor.runtime.triton_helpers import math as tl_math
import torch.nn as nn
import torch.nn.functional as F
from math import sqrt as sqrt
from itertools import product as product
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda


@triton.jit
def triton_poi_fused__softmax_add_log_mul_neg_0(in_ptr0, in_ptr1, out_ptr0,
    out_ptr1, xnumel, XBLOCK: tl.constexpr):
    xnumel = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + 4 * x0, xmask, eviction_policy='evict_last')
    tmp1 = tl.load(in_ptr1 + 4 * x0, xmask, eviction_policy='evict_last')
    tmp10 = tl.load(in_ptr0 + (1 + 4 * x0), xmask, eviction_policy='evict_last'
        )
    tmp11 = tl.load(in_ptr1 + (1 + 4 * x0), xmask, eviction_policy='evict_last'
        )
    tmp19 = tl.load(in_ptr0 + (2 + 4 * x0), xmask, eviction_policy='evict_last'
        )
    tmp20 = tl.load(in_ptr1 + (2 + 4 * x0), xmask, eviction_policy='evict_last'
        )
    tmp28 = tl.load(in_ptr0 + (3 + 4 * x0), xmask, eviction_policy='evict_last'
        )
    tmp29 = tl.load(in_ptr1 + (3 + 4 * x0), xmask, eviction_policy='evict_last'
        )
    tmp2 = tl_math.log(tmp1)
    tmp3 = -tmp2
    tmp4 = tl_math.log(tmp3)
    tmp5 = -0.001
    tmp6 = tmp4 * tmp5
    tmp7 = tmp0 + tmp6
    tmp8 = 1.0
    tmp9 = tmp7 * tmp8
    tmp12 = tl_math.log(tmp11)
    tmp13 = -tmp12
    tmp14 = tl_math.log(tmp13)
    tmp15 = tmp14 * tmp5
    tmp16 = tmp10 + tmp15
    tmp17 = tmp16 * tmp8
    tmp18 = triton_helpers.maximum(tmp9, tmp17)
    tmp21 = tl_math.log(tmp20)
    tmp22 = -tmp21
    tmp23 = tl_math.log(tmp22)
    tmp24 = tmp23 * tmp5
    tmp25 = tmp19 + tmp24
    tmp26 = tmp25 * tmp8
    tmp27 = triton_helpers.maximum(tmp18, tmp26)
    tmp30 = tl_math.log(tmp29)
    tmp31 = -tmp30
    tmp32 = tl_math.log(tmp31)
    tmp33 = tmp32 * tmp5
    tmp34 = tmp28 + tmp33
    tmp35 = tmp34 * tmp8
    tmp36 = triton_helpers.maximum(tmp27, tmp35)
    tmp37 = tmp9 - tmp36
    tmp38 = tmp37 * tmp8
    tmp39 = tl_math.exp(tmp38)
    tmp40 = tmp17 - tmp36
    tmp41 = tmp40 * tmp8
    tmp42 = tl_math.exp(tmp41)
    tmp43 = tmp39 + tmp42
    tmp44 = tmp26 - tmp36
    tmp45 = tmp44 * tmp8
    tmp46 = tl_math.exp(tmp45)
    tmp47 = tmp43 + tmp46
    tmp48 = tmp35 - tmp36
    tmp49 = tmp48 * tmp8
    tmp50 = tl_math.exp(tmp49)
    tmp51 = tmp47 + tmp50
    tl.store(out_ptr0 + x0, tmp36, xmask)
    tl.store(out_ptr1 + x0, tmp51, xmask)


@triton.jit
def triton_poi_fused__softmax_add_log_mul_neg_1(in_out_ptr0, in_ptr0,
    in_ptr1, in_ptr2, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp1 = tl.load(in_out_ptr0 + x2, xmask)
    tmp10 = tl.load(in_ptr1 + x1, xmask, eviction_policy='evict_last')
    tmp14 = tl.load(in_ptr2 + x1, xmask, eviction_policy='evict_last')
    tmp2 = tl_math.log(tmp1)
    tmp3 = -tmp2
    tmp4 = tl_math.log(tmp3)
    tmp5 = -0.001
    tmp6 = tmp4 * tmp5
    tmp7 = tmp0 + tmp6
    tmp8 = 1.0
    tmp9 = tmp7 * tmp8
    tmp11 = tmp9 - tmp10
    tmp12 = tmp11 * tmp8
    tmp13 = tl_math.exp(tmp12)
    tmp15 = tmp13 / tmp14
    tl.store(in_out_ptr0 + x2, tmp15, xmask)


@triton.jit
def triton_poi_fused_zeros_like_2(out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = 0.0
    tl.store(out_ptr0 + x0, tmp0, xmask)


@triton.jit
def triton_poi_fused_index_put_lift_fresh_zeros_like_3(in_ptr0, out_ptr0,
    xnumel, XBLOCK: tl.constexpr):
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    tl.full([XBLOCK], True, tl.int1)
    x1 = xindex // 64
    x0 = xindex % 64
    tmp0 = tl.load(in_ptr0 + x1, None, eviction_policy='evict_last')
    tmp1 = tl.full([XBLOCK], 4, tl.int32)
    tmp2 = tmp0 + tmp1
    tmp3 = tmp0 < 0
    tmp4 = tl.where(tmp3, tmp2, tmp0)
    tl.device_assert((0 <= tmp4) & (tmp4 < 4),
        'index out of bounds: 0 <= tmp4 < 4')
    tmp6 = 1.0
    tl.store(out_ptr0 + (x0 + 64 * tmp4), tmp6, None)


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)
        buf0 = torch.ops.aten.rand.default([4, 4, 4, 4], dtype=torch.
            float32, device=device(type='cuda', index=0), pin_memory=False)
        buf1 = buf0
        del buf0
        buf2 = empty_strided_cuda((4, 4, 4, 1), (16, 4, 1, 64), torch.float32)
        buf3 = empty_strided_cuda((4, 4, 4, 1), (16, 4, 1, 64), torch.float32)
        get_raw_stream(0)
        triton_poi_fused__softmax_add_log_mul_neg_0[grid(64)](arg0_1, buf1,
            buf2, buf3, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf4 = buf1
        del buf1
        triton_poi_fused__softmax_add_log_mul_neg_1[grid(256)](buf4, arg0_1,
            buf2, buf3, 256, XBLOCK=256, num_warps=4, num_stages=1)
        del arg0_1
        del buf2
        del buf3
        buf5 = torch.ops.aten.topk.default(buf4, 1)
        buf7 = buf5[1]
        del buf5
        buf8 = buf4
        del buf4
        triton_poi_fused_zeros_like_2[grid(256)](buf8, 256, XBLOCK=256,
            num_warps=4, num_stages=1)
        triton_poi_fused_index_put_lift_fresh_zeros_like_3[grid(4096)](buf7,
            buf8, 4096, XBLOCK=256, num_warps=4, num_stages=1)
        del buf7
    return buf8,


class _GumbelSoftMax(torch.autograd.Function):
    """
    implementing the MixedOp, but carried out in a different way as DARTS

    DARTS adds all operations together, then select the maximal one to construct the final network,
    however, during the late process, more weights are assigned to the None, this is unreasonable under the
    circumstance that per operation has the unsure number of inputs.

    Thus, we modifies the original DARTS by applying way in GDAS to test.

    This class aims to compute the gradients by ourself.
    """

    @staticmethod
    def forward(ctx, weights):
        weights_norm = F.softmax(weights, dim=-1)
        ctx.saved_for_backward = weights_norm
        mask = torch.zeros_like(weights_norm)
        _, idx = weights_norm.topk(dim=-1, k=1, largest=True)
        mask[idx] = 1.0
        return mask

    @staticmethod
    def backward(ctx, grad_output):
        gumbel_norm = ctx.saved_for_backward
        return gumbel_norm * (1 - gumbel_norm
            ) * grad_output * gumbel_norm.shape[0]


class GumbelSoftMaxNew(nn.Module):

    def __init__(self):
        super(GumbelSoftMaxNew, self).__init__()

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


Measured latency_ms: 0.071037