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 libdevice
from torch import nn
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda


@triton.jit
def triton_per_fused__to_copy_add_floor_lt_max_mean_mul_rsub_sigmoid_0(in_ptr0,
    in_ptr1, out_ptr1, xnumel, rnumel, XBLOCK: tl.constexpr):
    xnumel = 4
    RBLOCK: tl.constexpr = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r1 = rindex
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (r1 + 64 * x0), xmask, other=0.0)
    tmp1 = tl.load(in_ptr0 + (16 + r1 + 64 * x0), xmask, other=0.0)
    tmp3 = tl.load(in_ptr0 + (32 + r1 + 64 * x0), xmask, other=0.0)
    tmp5 = tl.load(in_ptr0 + (48 + r1 + 64 * x0), xmask, other=0.0)
    tmp13 = in_ptr1
    tmp2 = tmp0 + tmp1
    tmp4 = tmp2 + tmp3
    tmp6 = tmp4 + tmp5
    tmp7 = 4.0
    tmp8 = tmp6 / tmp7
    tmp9 = tl.broadcast_to(tmp8, [XBLOCK, RBLOCK])
    tmp11 = tl.where(xmask, tmp9, float('-inf'))
    tmp12 = triton_helpers.max2(tmp11, 1)[:, None]
    tmp14 = 0.8
    tmp15 = tmp13 + tmp14
    tmp16 = libdevice.floor(tmp15)
    tmp17 = 1.0
    tmp18 = tmp17 - tmp16
    tmp19 = tl.sigmoid(tmp8)
    tmp20 = tmp18 * tmp19
    tmp21 = 0.7
    tmp22 = tmp12 * tmp21
    tmp23 = tmp8 < tmp22
    tmp24 = tmp23.to(tl.float32)
    tmp25 = tmp16 * tmp24
    tmp26 = tmp20 + tmp25
    tl.store(out_ptr1 + (r1 + 16 * x0), tmp26, xmask)


@triton.jit
def triton_poi_fused__to_copy_add_div_floor_lt_mean_mul_rsub_sigmoid_1(in_ptr0,
    in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x3 = xindex
    x0 = xindex % 16
    x2 = xindex // 64
    tmp0 = tl.load(in_ptr0 + x3, xmask)
    tmp1 = tl.load(in_ptr1 + (x0 + 16 * x2), xmask, eviction_policy=
        'evict_last')
    tmp2 = tmp0 * tmp1
    tmp3 = tmp2 + tmp0
    tmp4 = 0.5
    tmp5 = tmp3 * tmp4
    tl.store(out_ptr0 + x3, tmp5, xmask)


def call(args):
    arg0_1, = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 4, 4), (64, 16, 4, 1))
    buf2 = torch.ops.aten.rand.default([], dtype=torch.float32, device=
        device(type='cpu'), pin_memory=False)
    buf3 = buf2
    del buf2
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf4 = empty_strided_cuda((4, 1, 4, 4), (16, 64, 4, 1), torch.float32)
        get_raw_stream(0)
        triton_per_fused__to_copy_add_floor_lt_max_mean_mul_rsub_sigmoid_0[grid
            (4)](arg0_1, buf3.item(), buf4, 4, 16, XBLOCK=1, num_warps=2,
            num_stages=1)
        del buf3
        buf5 = empty_strided_cuda((4, 4, 4, 4), (64, 16, 4, 1), torch.float32)
        triton_poi_fused__to_copy_add_div_floor_lt_mean_mul_rsub_sigmoid_1[grid
            (256)](arg0_1, buf4, buf5, 256, XBLOCK=128, num_warps=4,
            num_stages=1)
        del arg0_1
        del buf4
    return buf5,


class GDLNew(nn.Module):

    def __init__(self, drop_rate=0.8, drop_th=0.7):
        super(GDLNew, self).__init__()
        if not 0 <= drop_rate <= 1:
            raise ValueError('drop-rate must be in range [0, 1].')
        if not 0 <= drop_th <= 1:
            raise ValueError('drop-th must be in range [0, 1].')
        self.drop_rate = drop_rate
        self.drop_th = drop_th
        self.attention = None
        self.drop_mask = None

    def _select_map(self, importance_map, drop_mask):
        random_tensor = torch.rand([], dtype=torch.float32) + self.drop_rate
        binary_tensor = random_tensor.floor()
        return (1.0 - binary_tensor
            ) * importance_map + binary_tensor * drop_mask

    def _drop_mask(self, attention):
        b_size = attention.size(0)
        max_val, _ = torch.max(attention.view(b_size, -1), dim=1, keepdim=True)
        thr_val = max_val * self.drop_th
        thr_val = thr_val.view(b_size, 1, 1, 1)
        return (attention < thr_val).float()

    def extra_repr(self):
        return 'drop_rate={}, drop_th={}'.format(self.drop_rate, self.drop_th)

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


Measured latency_ms: 0.034313