# AOT ID: ['0_inference']
from ctypes import c_void_p, c_long, c_int
import torch
import math
import random
import os
import tempfile
from math import inf, nan
from torch._inductor.hooks import run_intermediate_hooks
from torch._inductor.utils import maybe_profile
from torch._inductor.codegen.memory_planning import _align as align
from torch import device, empty_strided
from torch._inductor.async_compile import AsyncCompile
from torch._inductor.select_algorithm import extern_kernels
from torch._inductor.codegen.multi_kernel import MultiKernelCall
import triton
import triton.language as tl
from torch._inductor.runtime.triton_heuristics import grid, split_scan_grid, grid_combo_kernels, start_graph, end_graph
from torch._C import _cuda_getCurrentRawStream as get_raw_stream

aten = torch.ops.aten
inductor_ops = torch.ops.inductor
_quantized = torch.ops._quantized
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cpu = torch._C._dynamo.guards._empty_strided_cpu
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
empty_strided_xpu = torch._C._dynamo.guards._empty_strided_xpu
reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor
alloc_from_pool = torch.ops.inductor._alloc_from_pool
async_compile = AsyncCompile()


# kernel path: runs/run_shard_7/inductor_cache/sw/cswbvd5l5fqejy3w27tz5i3b3hpstanzewjhsozvtttookjhk7zn.py
# Topologically Sorted Source Nodes: [log, neg, log_1, gumbel, add, weights_norm], Original ATen: [aten.log, aten.neg, aten.mul, aten.add, aten._softmax]
# Source node to ATen node mapping:
#   add => add
#   gumbel => mul
#   log => log
#   log_1 => log_1
#   neg => neg
#   weights_norm => exp, sum_1
# Graph fragment:
#   %log : [num_users=1] = call_function[target=torch.ops.aten.log.default](args = (%rand,), kwargs = {})
#   %neg : [num_users=1] = call_function[target=torch.ops.aten.neg.default](args = (%log,), kwargs = {})
#   %log_1 : [num_users=1] = call_function[target=torch.ops.aten.log.default](args = (%neg,), kwargs = {})
#   %mul : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%log_1, -0.001), kwargs = {})
#   %add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%arg0_1, %mul), kwargs = {})
#   %mul_tensor : [num_users=2] = call_function[target=torch.ops.aten.mul.Tensor](args = (%add, 1), kwargs = {})
#   %amax_default : [num_users=1] = call_function[target=torch.ops.aten.amax.default](args = (%mul_tensor, [-1], True), kwargs = {})
#   %sub_tensor : [num_users=1] = call_function[target=torch.ops.aten.sub.Tensor](args = (%mul_tensor, %amax_default), kwargs = {})
#   %div_tensor : [num_users=1] = call_function[target=torch.ops.aten.div.Tensor](args = (%sub_tensor, 1.0), kwargs = {})
#   %exp : [num_users=2] = call_function[target=torch.ops.aten.exp.default](args = (%div_tensor,), kwargs = {})
#   %sum_1 : [num_users=1] = call_function[target=torch.ops.aten.sum.dim_IntList](args = (%exp, [-1], True), kwargs = {})
triton_poi_fused__softmax_add_log_mul_neg_0 = async_compile.triton('triton_poi_fused__softmax_add_log_mul_neg_0', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, instance_descriptor, DeviceProperties

@triton_heuristics.pointwise(
    size_hints=[64], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: '*fp32', 3: '*fp32', 4: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2, 3, 4), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused__softmax_add_log_mul_neg_0', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 8, 'num_reduction': 0, 'backend_hash': 'A9C866B4A14FD3277824029365D703C2427B2E685E54EC9B3EF4ADC8D1EEAC1D', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
    min_elem_per_thread=0
)
@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)
''', device_str='cuda')


# kernel path: runs/run_shard_7/inductor_cache/t2/ct2shybnapdaeebefbjtzjs3x4xsmwbos5nap4ng2so72ngfdtot.py
# Topologically Sorted Source Nodes: [log, neg, log_1, gumbel, add, weights_norm], Original ATen: [aten.log, aten.neg, aten.mul, aten.add, aten._softmax]
# Source node to ATen node mapping:
#   add => add
#   gumbel => mul
#   log => log
#   log_1 => log_1
#   neg => neg
#   weights_norm => div_1, exp
# Graph fragment:
#   %log : [num_users=1] = call_function[target=torch.ops.aten.log.default](args = (%rand,), kwargs = {})
#   %neg : [num_users=1] = call_function[target=torch.ops.aten.neg.default](args = (%log,), kwargs = {})
#   %log_1 : [num_users=1] = call_function[target=torch.ops.aten.log.default](args = (%neg,), kwargs = {})
#   %mul : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%log_1, -0.001), kwargs = {})
#   %add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%arg0_1, %mul), kwargs = {})
#   %mul_tensor : [num_users=2] = call_function[target=torch.ops.aten.mul.Tensor](args = (%add, 1), kwargs = {})
#   %sub_tensor : [num_users=1] = call_function[target=torch.ops.aten.sub.Tensor](args = (%mul_tensor, %amax_default), kwargs = {})
#   %div_tensor : [num_users=1] = call_function[target=torch.ops.aten.div.Tensor](args = (%sub_tensor, 1.0), kwargs = {})
#   %exp : [num_users=2] = call_function[target=torch.ops.aten.exp.default](args = (%div_tensor,), kwargs = {})
#   %div_1 : [num_users=1] = call_function[target=torch.ops.aten.div.Tensor](args = (%exp, %sum_1), kwargs = {})
triton_poi_fused__softmax_add_log_mul_neg_1 = async_compile.triton('triton_poi_fused__softmax_add_log_mul_neg_1', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, instance_descriptor, DeviceProperties

@triton_heuristics.pointwise(
    size_hints=[256], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: '*fp32', 3: '*fp32', 4: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2, 3, 4), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused__softmax_add_log_mul_neg_1', 'mutated_arg_names': ['in_out_ptr0'], 'no_x_dim': False, 'num_load': 4, 'num_reduction': 0, 'backend_hash': 'A9C866B4A14FD3277824029365D703C2427B2E685E54EC9B3EF4ADC8D1EEAC1D', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
    min_elem_per_thread=0
)
@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)
''', device_str='cuda')


# kernel path: runs/run_shard_7/inductor_cache/zd/czdhovzy3i75dbb4lrrvs36jdqekoz77qi3q3aloxdafcrzuwdjh.py
# Topologically Sorted Source Nodes: [mask], Original ATen: [aten.zeros_like]
# Source node to ATen node mapping:
#   mask => full_default
# Graph fragment:
#   %full_default : [num_users=1] = call_function[target=torch.ops.aten.full.default](args = ([4, 4, 4, 4], 0), kwargs = {dtype: torch.float32, layout: torch.strided, device: cuda:0, pin_memory: False})
triton_poi_fused_zeros_like_2 = async_compile.triton('triton_poi_fused_zeros_like_2', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, instance_descriptor, DeviceProperties

@triton_heuristics.pointwise(
    size_hints=[256], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_zeros_like_2', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 0, 'num_reduction': 0, 'backend_hash': 'A9C866B4A14FD3277824029365D703C2427B2E685E54EC9B3EF4ADC8D1EEAC1D', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
    min_elem_per_thread=0
)
@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)
''', device_str='cuda')


# kernel path: runs/run_shard_7/inductor_cache/6t/c6t4l2n7ldxrmdakomrfhybejpedvrbykfwtgvtcdh32xkdnqc4p.py
# Topologically Sorted Source Nodes: [mask, setitem], Original ATen: [aten.zeros_like, aten.lift_fresh, aten.index_put]
# Source node to ATen node mapping:
#   mask => full_default
#   setitem => full_default_1, index_put
# Graph fragment:
#   %full_default : [num_users=1] = call_function[target=torch.ops.aten.full.default](args = ([4, 4, 4, 4], 0), kwargs = {dtype: torch.float32, layout: torch.strided, device: cuda:0, pin_memory: False})
#   %full_default_1 : [num_users=1] = call_function[target=torch.ops.aten.full.default](args = ([], 1.0), kwargs = {dtype: torch.float32, layout: torch.strided, device: cuda:0, pin_memory: False})
#   %index_put : [num_users=1] = call_function[target=torch.ops.aten.index_put_.default](args = (%full_default, [%getitem_1], %full_default_1), kwargs = {})
triton_poi_fused_index_put_lift_fresh_zeros_like_3 = async_compile.triton('triton_poi_fused_index_put_lift_fresh_zeros_like_3', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, instance_descriptor, DeviceProperties

@triton_heuristics.pointwise(
    size_hints=[4096], 
    filename=__file__,
    triton_meta={'signature': {0: '*i64', 1: '*fp32', 2: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_index_put_lift_fresh_zeros_like_3', 'mutated_arg_names': ['out_ptr0'], 'no_x_dim': False, 'num_load': 1, 'num_reduction': 0, 'backend_hash': 'A9C866B4A14FD3277824029365D703C2427B2E685E54EC9B3EF4ADC8D1EEAC1D', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
    min_elem_per_thread=0
)
@triton.jit
def triton_poi_fused_index_put_lift_fresh_zeros_like_3(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 4096
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = 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)
''', device_str='cuda')


async_compile.wait(globals())
del async_compile

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)
        # Topologically Sorted Source Nodes: [rand_like], Original ATen: [aten.rand_like]
        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)
        # Topologically Sorted Source Nodes: [log, neg, log_1, gumbel, add, weights_norm], Original ATen: [aten.log, aten.neg, aten.mul, aten.add, aten._softmax]
        stream0 = get_raw_stream(0)
        triton_poi_fused__softmax_add_log_mul_neg_0.run(arg0_1, buf1, buf2, buf3, 64, grid=grid(64), stream=stream0)
        buf4 = buf1; del buf1  # reuse
        # Topologically Sorted Source Nodes: [log, neg, log_1, gumbel, add, weights_norm], Original ATen: [aten.log, aten.neg, aten.mul, aten.add, aten._softmax]
        triton_poi_fused__softmax_add_log_mul_neg_1.run(buf4, arg0_1, buf2, buf3, 256, grid=grid(256), stream=stream0)
        del arg0_1
        del buf2
        del buf3
        # Topologically Sorted Source Nodes: [log, neg, log_1, gumbel, add, weights_norm, topk], Original ATen: [aten.log, aten.neg, aten.mul, aten.add, aten._softmax, aten.topk]
        buf5 = torch.ops.aten.topk.default(buf4, 1)
        buf7 = buf5[1]
        del buf5
        buf8 = buf4; del buf4  # reuse
        # Topologically Sorted Source Nodes: [mask], Original ATen: [aten.zeros_like]
        triton_poi_fused_zeros_like_2.run(buf8, 256, grid=grid(256), stream=stream0)
        # Topologically Sorted Source Nodes: [mask, setitem], Original ATen: [aten.zeros_like, aten.lift_fresh, aten.index_put]
        triton_poi_fused_index_put_lift_fresh_zeros_like_3.run(buf7, buf8, 4096, grid=grid(4096), stream=stream0)
        del buf7
    return (buf8, )


def benchmark_compiled_module(times=10, repeat=10):
    from torch._dynamo.testing import rand_strided
    from torch._inductor.utils import print_performance
    arg0_1 = rand_strided((4, 4, 4, 4), (64, 16, 4, 1), device='cuda:0', dtype=torch.float32)
    fn = lambda: call([arg0_1])
    return print_performance(fn, times=times, repeat=repeat)


if __name__ == "__main__":
    from torch._inductor.wrapper_benchmark import compiled_module_main
    compiled_module_main('None', benchmark_compiled_module)

import torch
import torch.nn as nn
import torch.nn.functional as F
from math import sqrt as sqrt
from itertools import product as product


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 GumbelSoftMax(nn.Module):

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

    def forward(self, weights, temp_coeff=1.0):
        gumbel = -0.001 * torch.log(-torch.log(torch.rand_like(weights)))
        weights = _GumbelSoftMax.apply((weights + gumbel) / temp_coeff)
        return weights


def get_inputs():
    return [torch.rand([4, 4, 4, 4])]


def get_init_inputs():
    return [[], {}]

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]
