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 import triton_helpers
import numpy as np
from torch import nn
import torch.nn.functional as F
from itertools import filterfalse
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda


@triton.jit
def _triton_helper_fn_add0(arg0_0, arg1_0):
    tmp0 = arg0_0 + arg1_0
    return tmp0


@triton.jit
def triton_per_fused_cumsum_index_mul_rsub_sort_sub_sum_0(in_ptr0, in_ptr1,
    out_ptr0, out_ptr2, out_ptr4, out_ptr5, out_ptr6, out_ptr7, out_ptr8,
    out_ptr9, xnumel, rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 64
    xoffset = tl.program_id(0) * XBLOCK
    xoffset + tl.arange(0, XBLOCK)[:, None]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r0 = rindex
    tmp0 = tl.load(in_ptr0 + r0, None)
    tmp1 = tl.load(in_ptr1 + r0, None)
    tmp2 = 2.0
    tmp3 = tmp1 * tmp2
    tmp4 = 1.0
    tmp5 = tmp3 - tmp4
    tmp6 = tmp0 * tmp5
    tmp7 = tmp4 - tmp6
    tmp8 = r0
    tmp9 = tmp8.to(tl.int16)
    tmp10 = tl.broadcast_to(tmp7, [XBLOCK, RBLOCK])
    tmp11 = tl.broadcast_to(tmp9, [XBLOCK, RBLOCK])
    tmp12, tmp13 = triton_helpers.sort_with_index(tmp10, tmp11, None, 1,
        stable=False, descending=True)
    tmp14 = -tmp0
    tmp15 = tmp4 - tmp1
    tmp16 = tmp15 * tmp2
    tmp17 = tmp16 - tmp4
    tmp18 = tmp14 * tmp17
    tmp19 = tmp4 - tmp18
    tmp20 = tl.broadcast_to(tmp19, [XBLOCK, RBLOCK])
    tmp21, tmp22 = triton_helpers.sort_with_index(tmp20, tmp11, None, 1,
        stable=False, descending=True)
    tmp23 = tmp13.to(tl.int64)
    tmp24 = tl.full([XBLOCK, RBLOCK], 64, tl.int32)
    tmp25 = tmp23 + tmp24
    tmp26 = tmp23 < 0
    tmp27 = tl.where(tmp26, tmp25, tmp23)
    tl.device_assert((0 <= tmp27) & (tmp27 < 64),
        'index out of bounds: 0 <= tmp27 < 64')
    tmp29 = tl.load(in_ptr1 + tmp27, None, eviction_policy='evict_last')
    tmp30 = tl.broadcast_to(tmp29, [XBLOCK, RBLOCK])
    tmp32 = tl.sum(tmp30, 1)[:, None]
    tmp33 = tmp29.to(tl.float32)
    tmp34 = tl.broadcast_to(tmp33, [XBLOCK, RBLOCK])
    tmp35, = tl.associative_scan((tmp34,), 1, _triton_helper_fn_add0)
    tmp36 = tmp4 - tmp29
    tmp37 = tmp36.to(tl.float32)
    tmp38 = tl.broadcast_to(tmp37, [XBLOCK, RBLOCK])
    tmp39, = tl.associative_scan((tmp38,), 1, _triton_helper_fn_add0)
    tmp40 = tmp22.to(tl.int64)
    tmp41 = tmp40 + tmp24
    tmp42 = tmp40 < 0
    tmp43 = tl.where(tmp42, tmp41, tmp40)
    tl.device_assert((0 <= tmp43) & (tmp43 < 64),
        'index out of bounds: 0 <= tmp43 < 64')
    tmp45 = tl.load(in_ptr1 + tmp43 % 64, None, eviction_policy='evict_last')
    tmp46 = tmp4 - tmp45
    tmp47 = tl.broadcast_to(tmp46, [XBLOCK, RBLOCK])
    tmp49 = tl.sum(tmp47, 1)[:, None]
    tmp50 = tmp46.to(tl.float32)
    tmp51 = tl.broadcast_to(tmp50, [XBLOCK, RBLOCK])
    tmp52, = tl.associative_scan((tmp51,), 1, _triton_helper_fn_add0)
    tmp53 = tmp4 - tmp46
    tmp54 = tmp53.to(tl.float32)
    tmp55 = tl.broadcast_to(tmp54, [XBLOCK, RBLOCK])
    tmp56, = tl.associative_scan((tmp55,), 1, _triton_helper_fn_add0)
    tl.store(out_ptr0 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp12, None)
    tl.store(out_ptr2 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp21, None)
    tl.store(out_ptr5 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp35, None)
    tl.store(out_ptr6 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp39, None)
    tl.store(out_ptr8 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp52, None)
    tl.store(out_ptr9 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp56, None)
    tl.store(out_ptr4 + tl.full([XBLOCK, 1], 0, tl.int32), tmp32, None)
    tl.store(out_ptr7 + tl.full([XBLOCK, 1], 0, tl.int32), tmp49, None)


@triton.jit
def _triton_helper_fn_add0(arg0_0, arg1_0):
    tmp0 = arg0_0 + arg1_0
    return tmp0


@triton.jit
def triton_per_fused_cumsum_index_mul_rsub_sort_sub_sum_1(in_ptr0, in_ptr1,
    out_ptr0, out_ptr2, out_ptr4, out_ptr5, out_ptr6, out_ptr7, out_ptr8,
    out_ptr9, xnumel, rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 64
    xoffset = tl.program_id(0) * XBLOCK
    xoffset + tl.arange(0, XBLOCK)[:, None]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r0 = rindex
    tmp0 = tl.load(in_ptr0 + (64 + r0), None)
    tmp1 = tl.load(in_ptr1 + (64 + r0), None)
    tmp2 = 2.0
    tmp3 = tmp1 * tmp2
    tmp4 = 1.0
    tmp5 = tmp3 - tmp4
    tmp6 = tmp0 * tmp5
    tmp7 = tmp4 - tmp6
    tmp8 = r0
    tmp9 = tmp8.to(tl.int16)
    tmp10 = tl.broadcast_to(tmp7, [XBLOCK, RBLOCK])
    tmp11 = tl.broadcast_to(tmp9, [XBLOCK, RBLOCK])
    tmp12, tmp13 = triton_helpers.sort_with_index(tmp10, tmp11, None, 1,
        stable=False, descending=True)
    tmp14 = -tmp0
    tmp15 = tmp4 - tmp1
    tmp16 = tmp15 * tmp2
    tmp17 = tmp16 - tmp4
    tmp18 = tmp14 * tmp17
    tmp19 = tmp4 - tmp18
    tmp20 = tl.broadcast_to(tmp19, [XBLOCK, RBLOCK])
    tmp21, tmp22 = triton_helpers.sort_with_index(tmp20, tmp11, None, 1,
        stable=False, descending=True)
    tmp23 = tmp13.to(tl.int64)
    tmp24 = tl.full([XBLOCK, RBLOCK], 64, tl.int32)
    tmp25 = tmp23 + tmp24
    tmp26 = tmp23 < 0
    tmp27 = tl.where(tmp26, tmp25, tmp23)
    tl.device_assert((0 <= tmp27) & (tmp27 < 64),
        'index out of bounds: 0 <= tmp27 < 64')
    tmp29 = tl.load(in_ptr1 + (64 + tmp27), None, eviction_policy='evict_last')
    tmp30 = tl.broadcast_to(tmp29, [XBLOCK, RBLOCK])
    tmp32 = tl.sum(tmp30, 1)[:, None]
    tmp33 = tmp29.to(tl.float32)
    tmp34 = tl.broadcast_to(tmp33, [XBLOCK, RBLOCK])
    tmp35, = tl.associative_scan((tmp34,), 1, _triton_helper_fn_add0)
    tmp36 = tmp4 - tmp29
    tmp37 = tmp36.to(tl.float32)
    tmp38 = tl.broadcast_to(tmp37, [XBLOCK, RBLOCK])
    tmp39, = tl.associative_scan((tmp38,), 1, _triton_helper_fn_add0)
    tmp40 = tmp22.to(tl.int64)
    tmp41 = tmp40 + tmp24
    tmp42 = tmp40 < 0
    tmp43 = tl.where(tmp42, tmp41, tmp40)
    tl.device_assert((0 <= tmp43) & (tmp43 < 64),
        'index out of bounds: 0 <= tmp43 < 64')
    tmp45 = tl.load(in_ptr1 + (64 + tmp43 % 64), None, eviction_policy=
        'evict_last')
    tmp46 = tmp4 - tmp45
    tmp47 = tl.broadcast_to(tmp46, [XBLOCK, RBLOCK])
    tmp49 = tl.sum(tmp47, 1)[:, None]
    tmp50 = tmp46.to(tl.float32)
    tmp51 = tl.broadcast_to(tmp50, [XBLOCK, RBLOCK])
    tmp52, = tl.associative_scan((tmp51,), 1, _triton_helper_fn_add0)
    tmp53 = tmp4 - tmp46
    tmp54 = tmp53.to(tl.float32)
    tmp55 = tl.broadcast_to(tmp54, [XBLOCK, RBLOCK])
    tmp56, = tl.associative_scan((tmp55,), 1, _triton_helper_fn_add0)
    tl.store(out_ptr0 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp12, None)
    tl.store(out_ptr2 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp21, None)
    tl.store(out_ptr5 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp35, None)
    tl.store(out_ptr6 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp39, None)
    tl.store(out_ptr8 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp52, None)
    tl.store(out_ptr9 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp56, None)
    tl.store(out_ptr4 + tl.full([XBLOCK, 1], 0, tl.int32), tmp32, None)
    tl.store(out_ptr7 + tl.full([XBLOCK, 1], 0, tl.int32), tmp49, None)


@triton.jit
def _triton_helper_fn_add0(arg0_0, arg1_0):
    tmp0 = arg0_0 + arg1_0
    return tmp0


@triton.jit
def triton_per_fused_cumsum_index_mul_rsub_sort_sub_sum_2(in_ptr0, in_ptr1,
    out_ptr0, out_ptr2, out_ptr4, out_ptr5, out_ptr6, out_ptr7, out_ptr8,
    out_ptr9, xnumel, rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 64
    xoffset = tl.program_id(0) * XBLOCK
    xoffset + tl.arange(0, XBLOCK)[:, None]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r0 = rindex
    tmp0 = tl.load(in_ptr0 + (128 + r0), None)
    tmp1 = tl.load(in_ptr1 + (128 + r0), None)
    tmp2 = 2.0
    tmp3 = tmp1 * tmp2
    tmp4 = 1.0
    tmp5 = tmp3 - tmp4
    tmp6 = tmp0 * tmp5
    tmp7 = tmp4 - tmp6
    tmp8 = r0
    tmp9 = tmp8.to(tl.int16)
    tmp10 = tl.broadcast_to(tmp7, [XBLOCK, RBLOCK])
    tmp11 = tl.broadcast_to(tmp9, [XBLOCK, RBLOCK])
    tmp12, tmp13 = triton_helpers.sort_with_index(tmp10, tmp11, None, 1,
        stable=False, descending=True)
    tmp14 = -tmp0
    tmp15 = tmp4 - tmp1
    tmp16 = tmp15 * tmp2
    tmp17 = tmp16 - tmp4
    tmp18 = tmp14 * tmp17
    tmp19 = tmp4 - tmp18
    tmp20 = tl.broadcast_to(tmp19, [XBLOCK, RBLOCK])
    tmp21, tmp22 = triton_helpers.sort_with_index(tmp20, tmp11, None, 1,
        stable=False, descending=True)
    tmp23 = tmp13.to(tl.int64)
    tmp24 = tl.full([XBLOCK, RBLOCK], 64, tl.int32)
    tmp25 = tmp23 + tmp24
    tmp26 = tmp23 < 0
    tmp27 = tl.where(tmp26, tmp25, tmp23)
    tl.device_assert((0 <= tmp27) & (tmp27 < 64),
        'index out of bounds: 0 <= tmp27 < 64')
    tmp29 = tl.load(in_ptr1 + (128 + tmp27), None, eviction_policy='evict_last'
        )
    tmp30 = tl.broadcast_to(tmp29, [XBLOCK, RBLOCK])
    tmp32 = tl.sum(tmp30, 1)[:, None]
    tmp33 = tmp29.to(tl.float32)
    tmp34 = tl.broadcast_to(tmp33, [XBLOCK, RBLOCK])
    tmp35, = tl.associative_scan((tmp34,), 1, _triton_helper_fn_add0)
    tmp36 = tmp4 - tmp29
    tmp37 = tmp36.to(tl.float32)
    tmp38 = tl.broadcast_to(tmp37, [XBLOCK, RBLOCK])
    tmp39, = tl.associative_scan((tmp38,), 1, _triton_helper_fn_add0)
    tmp40 = tmp22.to(tl.int64)
    tmp41 = tmp40 + tmp24
    tmp42 = tmp40 < 0
    tmp43 = tl.where(tmp42, tmp41, tmp40)
    tl.device_assert((0 <= tmp43) & (tmp43 < 64),
        'index out of bounds: 0 <= tmp43 < 64')
    tmp45 = tl.load(in_ptr1 + (128 + tmp43 % 64), None, eviction_policy=
        'evict_last')
    tmp46 = tmp4 - tmp45
    tmp47 = tl.broadcast_to(tmp46, [XBLOCK, RBLOCK])
    tmp49 = tl.sum(tmp47, 1)[:, None]
    tmp50 = tmp46.to(tl.float32)
    tmp51 = tl.broadcast_to(tmp50, [XBLOCK, RBLOCK])
    tmp52, = tl.associative_scan((tmp51,), 1, _triton_helper_fn_add0)
    tmp53 = tmp4 - tmp46
    tmp54 = tmp53.to(tl.float32)
    tmp55 = tl.broadcast_to(tmp54, [XBLOCK, RBLOCK])
    tmp56, = tl.associative_scan((tmp55,), 1, _triton_helper_fn_add0)
    tl.store(out_ptr0 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp12, None)
    tl.store(out_ptr2 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp21, None)
    tl.store(out_ptr5 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp35, None)
    tl.store(out_ptr6 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp39, None)
    tl.store(out_ptr8 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp52, None)
    tl.store(out_ptr9 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp56, None)
    tl.store(out_ptr4 + tl.full([XBLOCK, 1], 0, tl.int32), tmp32, None)
    tl.store(out_ptr7 + tl.full([XBLOCK, 1], 0, tl.int32), tmp49, None)


@triton.jit
def _triton_helper_fn_add0(arg0_0, arg1_0):
    tmp0 = arg0_0 + arg1_0
    return tmp0


@triton.jit
def triton_per_fused_cumsum_index_mul_rsub_sort_sub_sum_3(in_ptr0, in_ptr1,
    out_ptr0, out_ptr2, out_ptr4, out_ptr5, out_ptr6, out_ptr7, out_ptr8,
    out_ptr9, xnumel, rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 64
    xoffset = tl.program_id(0) * XBLOCK
    xoffset + tl.arange(0, XBLOCK)[:, None]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r0 = rindex
    tmp0 = tl.load(in_ptr0 + (192 + r0), None)
    tmp1 = tl.load(in_ptr1 + (192 + r0), None)
    tmp2 = 2.0
    tmp3 = tmp1 * tmp2
    tmp4 = 1.0
    tmp5 = tmp3 - tmp4
    tmp6 = tmp0 * tmp5
    tmp7 = tmp4 - tmp6
    tmp8 = r0
    tmp9 = tmp8.to(tl.int16)
    tmp10 = tl.broadcast_to(tmp7, [XBLOCK, RBLOCK])
    tmp11 = tl.broadcast_to(tmp9, [XBLOCK, RBLOCK])
    tmp12, tmp13 = triton_helpers.sort_with_index(tmp10, tmp11, None, 1,
        stable=False, descending=True)
    tmp14 = -tmp0
    tmp15 = tmp4 - tmp1
    tmp16 = tmp15 * tmp2
    tmp17 = tmp16 - tmp4
    tmp18 = tmp14 * tmp17
    tmp19 = tmp4 - tmp18
    tmp20 = tl.broadcast_to(tmp19, [XBLOCK, RBLOCK])
    tmp21, tmp22 = triton_helpers.sort_with_index(tmp20, tmp11, None, 1,
        stable=False, descending=True)
    tmp23 = tmp13.to(tl.int64)
    tmp24 = tl.full([XBLOCK, RBLOCK], 64, tl.int32)
    tmp25 = tmp23 + tmp24
    tmp26 = tmp23 < 0
    tmp27 = tl.where(tmp26, tmp25, tmp23)
    tl.device_assert((0 <= tmp27) & (tmp27 < 64),
        'index out of bounds: 0 <= tmp27 < 64')
    tmp29 = tl.load(in_ptr1 + (192 + tmp27), None, eviction_policy='evict_last'
        )
    tmp30 = tl.broadcast_to(tmp29, [XBLOCK, RBLOCK])
    tmp32 = tl.sum(tmp30, 1)[:, None]
    tmp33 = tmp29.to(tl.float32)
    tmp34 = tl.broadcast_to(tmp33, [XBLOCK, RBLOCK])
    tmp35, = tl.associative_scan((tmp34,), 1, _triton_helper_fn_add0)
    tmp36 = tmp4 - tmp29
    tmp37 = tmp36.to(tl.float32)
    tmp38 = tl.broadcast_to(tmp37, [XBLOCK, RBLOCK])
    tmp39, = tl.associative_scan((tmp38,), 1, _triton_helper_fn_add0)
    tmp40 = tmp22.to(tl.int64)
    tmp41 = tmp40 + tmp24
    tmp42 = tmp40 < 0
    tmp43 = tl.where(tmp42, tmp41, tmp40)
    tl.device_assert((0 <= tmp43) & (tmp43 < 64),
        'index out of bounds: 0 <= tmp43 < 64')
    tmp45 = tl.load(in_ptr1 + (192 + tmp43 % 64), None, eviction_policy=
        'evict_last')
    tmp46 = tmp4 - tmp45
    tmp47 = tl.broadcast_to(tmp46, [XBLOCK, RBLOCK])
    tmp49 = tl.sum(tmp47, 1)[:, None]
    tmp50 = tmp46.to(tl.float32)
    tmp51 = tl.broadcast_to(tmp50, [XBLOCK, RBLOCK])
    tmp52, = tl.associative_scan((tmp51,), 1, _triton_helper_fn_add0)
    tmp53 = tmp4 - tmp46
    tmp54 = tmp53.to(tl.float32)
    tmp55 = tl.broadcast_to(tmp54, [XBLOCK, RBLOCK])
    tmp56, = tl.associative_scan((tmp55,), 1, _triton_helper_fn_add0)
    tl.store(out_ptr0 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp12, None)
    tl.store(out_ptr2 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp21, None)
    tl.store(out_ptr5 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp35, None)
    tl.store(out_ptr6 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp39, None)
    tl.store(out_ptr8 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp52, None)
    tl.store(out_ptr9 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp56, None)
    tl.store(out_ptr4 + tl.full([XBLOCK, 1], 0, tl.int32), tmp32, None)
    tl.store(out_ptr7 + tl.full([XBLOCK, 1], 0, tl.int32), tmp49, None)


@triton.jit
def triton_per_fused_add_div_dot_relu_rsub_sub_4(in_out_ptr0, in_out_ptr1,
    in_out_ptr2, in_out_ptr3, in_out_ptr4, in_out_ptr5, in_out_ptr6,
    in_out_ptr7, in_ptr0, in_ptr1, in_ptr2, in_ptr3, in_ptr4, in_ptr5,
    in_ptr6, in_ptr7, in_ptr8, in_ptr9, in_ptr10, in_ptr11, in_ptr12,
    in_ptr13, in_ptr14, in_ptr15, in_ptr16, in_ptr17, in_ptr18, in_ptr19,
    in_ptr20, in_ptr21, in_ptr22, in_ptr23, xnumel, rnumel, XBLOCK: tl.
    constexpr):
    RBLOCK: tl.constexpr = 64
    xoffset = tl.program_id(0) * XBLOCK
    xoffset + tl.arange(0, XBLOCK)[:, None]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r0 = rindex
    tmp0 = tl.load(in_ptr0 + r0, None)
    tmp6 = tl.load(in_out_ptr0 + 0)
    tmp7 = tl.broadcast_to(tmp6, [XBLOCK, RBLOCK])
    tmp24 = tl.load(in_ptr1 + r0, None)
    tmp26 = tl.load(in_ptr2 + r0, None)
    tmp35 = tl.load(in_ptr3 + r0, None)
    tmp37 = tl.load(in_out_ptr1 + 0)
    tmp38 = tl.broadcast_to(tmp37, [XBLOCK, RBLOCK])
    tmp54 = tl.load(in_ptr4 + r0, None)
    tmp56 = tl.load(in_ptr5 + r0, None)
    tmp65 = tl.load(in_ptr6 + r0, None)
    tmp67 = tl.load(in_out_ptr2 + 0)
    tmp68 = tl.broadcast_to(tmp67, [XBLOCK, RBLOCK])
    tmp84 = tl.load(in_ptr7 + r0, None)
    tmp86 = tl.load(in_ptr8 + r0, None)
    tmp95 = tl.load(in_ptr9 + r0, None)
    tmp97 = tl.load(in_out_ptr3 + 0)
    tmp98 = tl.broadcast_to(tmp97, [XBLOCK, RBLOCK])
    tmp114 = tl.load(in_ptr10 + r0, None)
    tmp116 = tl.load(in_ptr11 + r0, None)
    tmp125 = tl.load(in_ptr12 + r0, None)
    tmp127 = tl.load(in_out_ptr4 + 0)
    tmp128 = tl.broadcast_to(tmp127, [XBLOCK, RBLOCK])
    tmp144 = tl.load(in_ptr13 + r0, None)
    tmp146 = tl.load(in_ptr14 + r0, None)
    tmp155 = tl.load(in_ptr15 + r0, None)
    tmp157 = tl.load(in_out_ptr5 + 0)
    tmp158 = tl.broadcast_to(tmp157, [XBLOCK, RBLOCK])
    tmp174 = tl.load(in_ptr16 + r0, None)
    tmp176 = tl.load(in_ptr17 + r0, None)
    tmp185 = tl.load(in_ptr18 + r0, None)
    tmp187 = tl.load(in_out_ptr6 + 0)
    tmp188 = tl.broadcast_to(tmp187, [XBLOCK, RBLOCK])
    tmp204 = tl.load(in_ptr19 + r0, None)
    tmp206 = tl.load(in_ptr20 + r0, None)
    tmp215 = tl.load(in_ptr21 + r0, None)
    tmp217 = tl.load(in_out_ptr7 + 0)
    tmp218 = tl.broadcast_to(tmp217, [XBLOCK, RBLOCK])
    tmp234 = tl.load(in_ptr22 + r0, None)
    tmp236 = tl.load(in_ptr23 + r0, None)
    tmp1 = tl.full([1, 1], 0, tl.int32)
    tmp2 = triton_helpers.maximum(tmp1, tmp0)
    tmp3 = r0
    tmp4 = tl.full([1, 1], 1, tl.int64)
    tmp5 = tmp3 >= tmp4
    tmp8 = tl.load(in_ptr1 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp5,
        other=0.0)
    tmp9 = tmp7 - tmp8
    tmp10 = tl.load(in_ptr2 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp5,
        other=0.0)
    tmp11 = tmp7 + tmp10
    tmp12 = tmp9 / tmp11
    tmp13 = 1.0
    tmp14 = tmp13 - tmp12
    tmp15 = tl.load(in_ptr1 + tl.broadcast_to(-1 + r0, [XBLOCK, RBLOCK]),
        tmp5, other=0.0)
    tmp16 = tmp7 - tmp15
    tmp17 = tl.load(in_ptr2 + tl.broadcast_to(-1 + r0, [XBLOCK, RBLOCK]),
        tmp5, other=0.0)
    tmp18 = tmp7 + tmp17
    tmp19 = tmp16 / tmp18
    tmp20 = tmp13 - tmp19
    tmp21 = tmp14 - tmp20
    tmp22 = tl.full(tmp21.shape, 0.0, tmp21.dtype)
    tmp23 = tl.where(tmp5, tmp21, tmp22)
    tmp25 = tmp7 - tmp24
    tmp27 = tmp7 + tmp26
    tmp28 = tmp25 / tmp27
    tmp29 = tmp13 - tmp28
    tmp30 = tl.where(tmp5, tmp23, tmp29)
    tmp31 = tmp2 * tmp30
    tmp32 = tl.broadcast_to(tmp31, [XBLOCK, RBLOCK])
    tmp34 = tl.sum(tmp32, 1)[:, None]
    tmp36 = triton_helpers.maximum(tmp1, tmp35)
    tmp39 = tl.load(in_ptr4 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp5,
        other=0.0)
    tmp40 = tmp38 - tmp39
    tmp41 = tl.load(in_ptr5 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp5,
        other=0.0)
    tmp42 = tmp38 + tmp41
    tmp43 = tmp40 / tmp42
    tmp44 = tmp13 - tmp43
    tmp45 = tl.load(in_ptr4 + tl.broadcast_to(-1 + r0, [XBLOCK, RBLOCK]),
        tmp5, other=0.0)
    tmp46 = tmp38 - tmp45
    tmp47 = tl.load(in_ptr5 + tl.broadcast_to(-1 + r0, [XBLOCK, RBLOCK]),
        tmp5, other=0.0)
    tmp48 = tmp38 + tmp47
    tmp49 = tmp46 / tmp48
    tmp50 = tmp13 - tmp49
    tmp51 = tmp44 - tmp50
    tmp52 = tl.full(tmp51.shape, 0.0, tmp51.dtype)
    tmp53 = tl.where(tmp5, tmp51, tmp52)
    tmp55 = tmp38 - tmp54
    tmp57 = tmp38 + tmp56
    tmp58 = tmp55 / tmp57
    tmp59 = tmp13 - tmp58
    tmp60 = tl.where(tmp5, tmp53, tmp59)
    tmp61 = tmp36 * tmp60
    tmp62 = tl.broadcast_to(tmp61, [XBLOCK, RBLOCK])
    tmp64 = tl.sum(tmp62, 1)[:, None]
    tmp66 = triton_helpers.maximum(tmp1, tmp65)
    tmp69 = tl.load(in_ptr7 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp5,
        other=0.0)
    tmp70 = tmp68 - tmp69
    tmp71 = tl.load(in_ptr8 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp5,
        other=0.0)
    tmp72 = tmp68 + tmp71
    tmp73 = tmp70 / tmp72
    tmp74 = tmp13 - tmp73
    tmp75 = tl.load(in_ptr7 + tl.broadcast_to(-1 + r0, [XBLOCK, RBLOCK]),
        tmp5, other=0.0)
    tmp76 = tmp68 - tmp75
    tmp77 = tl.load(in_ptr8 + tl.broadcast_to(-1 + r0, [XBLOCK, RBLOCK]),
        tmp5, other=0.0)
    tmp78 = tmp68 + tmp77
    tmp79 = tmp76 / tmp78
    tmp80 = tmp13 - tmp79
    tmp81 = tmp74 - tmp80
    tmp82 = tl.full(tmp81.shape, 0.0, tmp81.dtype)
    tmp83 = tl.where(tmp5, tmp81, tmp82)
    tmp85 = tmp68 - tmp84
    tmp87 = tmp68 + tmp86
    tmp88 = tmp85 / tmp87
    tmp89 = tmp13 - tmp88
    tmp90 = tl.where(tmp5, tmp83, tmp89)
    tmp91 = tmp66 * tmp90
    tmp92 = tl.broadcast_to(tmp91, [XBLOCK, RBLOCK])
    tmp94 = tl.sum(tmp92, 1)[:, None]
    tmp96 = triton_helpers.maximum(tmp1, tmp95)
    tmp99 = tl.load(in_ptr10 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp5,
        other=0.0)
    tmp100 = tmp98 - tmp99
    tmp101 = tl.load(in_ptr11 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp5,
        other=0.0)
    tmp102 = tmp98 + tmp101
    tmp103 = tmp100 / tmp102
    tmp104 = tmp13 - tmp103
    tmp105 = tl.load(in_ptr10 + tl.broadcast_to(-1 + r0, [XBLOCK, RBLOCK]),
        tmp5, other=0.0)
    tmp106 = tmp98 - tmp105
    tmp107 = tl.load(in_ptr11 + tl.broadcast_to(-1 + r0, [XBLOCK, RBLOCK]),
        tmp5, other=0.0)
    tmp108 = tmp98 + tmp107
    tmp109 = tmp106 / tmp108
    tmp110 = tmp13 - tmp109
    tmp111 = tmp104 - tmp110
    tmp112 = tl.full(tmp111.shape, 0.0, tmp111.dtype)
    tmp113 = tl.where(tmp5, tmp111, tmp112)
    tmp115 = tmp98 - tmp114
    tmp117 = tmp98 + tmp116
    tmp118 = tmp115 / tmp117
    tmp119 = tmp13 - tmp118
    tmp120 = tl.where(tmp5, tmp113, tmp119)
    tmp121 = tmp96 * tmp120
    tmp122 = tl.broadcast_to(tmp121, [XBLOCK, RBLOCK])
    tmp124 = tl.sum(tmp122, 1)[:, None]
    tmp126 = triton_helpers.maximum(tmp1, tmp125)
    tmp129 = tl.load(in_ptr13 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp5,
        other=0.0)
    tmp130 = tmp128 - tmp129
    tmp131 = tl.load(in_ptr14 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp5,
        other=0.0)
    tmp132 = tmp128 + tmp131
    tmp133 = tmp130 / tmp132
    tmp134 = tmp13 - tmp133
    tmp135 = tl.load(in_ptr13 + tl.broadcast_to(-1 + r0, [XBLOCK, RBLOCK]),
        tmp5, other=0.0)
    tmp136 = tmp128 - tmp135
    tmp137 = tl.load(in_ptr14 + tl.broadcast_to(-1 + r0, [XBLOCK, RBLOCK]),
        tmp5, other=0.0)
    tmp138 = tmp128 + tmp137
    tmp139 = tmp136 / tmp138
    tmp140 = tmp13 - tmp139
    tmp141 = tmp134 - tmp140
    tmp142 = tl.full(tmp141.shape, 0.0, tmp141.dtype)
    tmp143 = tl.where(tmp5, tmp141, tmp142)
    tmp145 = tmp128 - tmp144
    tmp147 = tmp128 + tmp146
    tmp148 = tmp145 / tmp147
    tmp149 = tmp13 - tmp148
    tmp150 = tl.where(tmp5, tmp143, tmp149)
    tmp151 = tmp126 * tmp150
    tmp152 = tl.broadcast_to(tmp151, [XBLOCK, RBLOCK])
    tmp154 = tl.sum(tmp152, 1)[:, None]
    tmp156 = triton_helpers.maximum(tmp1, tmp155)
    tmp159 = tl.load(in_ptr16 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp5,
        other=0.0)
    tmp160 = tmp158 - tmp159
    tmp161 = tl.load(in_ptr17 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp5,
        other=0.0)
    tmp162 = tmp158 + tmp161
    tmp163 = tmp160 / tmp162
    tmp164 = tmp13 - tmp163
    tmp165 = tl.load(in_ptr16 + tl.broadcast_to(-1 + r0, [XBLOCK, RBLOCK]),
        tmp5, other=0.0)
    tmp166 = tmp158 - tmp165
    tmp167 = tl.load(in_ptr17 + tl.broadcast_to(-1 + r0, [XBLOCK, RBLOCK]),
        tmp5, other=0.0)
    tmp168 = tmp158 + tmp167
    tmp169 = tmp166 / tmp168
    tmp170 = tmp13 - tmp169
    tmp171 = tmp164 - tmp170
    tmp172 = tl.full(tmp171.shape, 0.0, tmp171.dtype)
    tmp173 = tl.where(tmp5, tmp171, tmp172)
    tmp175 = tmp158 - tmp174
    tmp177 = tmp158 + tmp176
    tmp178 = tmp175 / tmp177
    tmp179 = tmp13 - tmp178
    tmp180 = tl.where(tmp5, tmp173, tmp179)
    tmp181 = tmp156 * tmp180
    tmp182 = tl.broadcast_to(tmp181, [XBLOCK, RBLOCK])
    tmp184 = tl.sum(tmp182, 1)[:, None]
    tmp186 = triton_helpers.maximum(tmp1, tmp185)
    tmp189 = tl.load(in_ptr19 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp5,
        other=0.0)
    tmp190 = tmp188 - tmp189
    tmp191 = tl.load(in_ptr20 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp5,
        other=0.0)
    tmp192 = tmp188 + tmp191
    tmp193 = tmp190 / tmp192
    tmp194 = tmp13 - tmp193
    tmp195 = tl.load(in_ptr19 + tl.broadcast_to(-1 + r0, [XBLOCK, RBLOCK]),
        tmp5, other=0.0)
    tmp196 = tmp188 - tmp195
    tmp197 = tl.load(in_ptr20 + tl.broadcast_to(-1 + r0, [XBLOCK, RBLOCK]),
        tmp5, other=0.0)
    tmp198 = tmp188 + tmp197
    tmp199 = tmp196 / tmp198
    tmp200 = tmp13 - tmp199
    tmp201 = tmp194 - tmp200
    tmp202 = tl.full(tmp201.shape, 0.0, tmp201.dtype)
    tmp203 = tl.where(tmp5, tmp201, tmp202)
    tmp205 = tmp188 - tmp204
    tmp207 = tmp188 + tmp206
    tmp208 = tmp205 / tmp207
    tmp209 = tmp13 - tmp208
    tmp210 = tl.where(tmp5, tmp203, tmp209)
    tmp211 = tmp186 * tmp210
    tmp212 = tl.broadcast_to(tmp211, [XBLOCK, RBLOCK])
    tmp214 = tl.sum(tmp212, 1)[:, None]
    tmp216 = triton_helpers.maximum(tmp1, tmp215)
    tmp219 = tl.load(in_ptr22 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp5,
        other=0.0)
    tmp220 = tmp218 - tmp219
    tmp221 = tl.load(in_ptr23 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp5,
        other=0.0)
    tmp222 = tmp218 + tmp221
    tmp223 = tmp220 / tmp222
    tmp224 = tmp13 - tmp223
    tmp225 = tl.load(in_ptr22 + tl.broadcast_to(-1 + r0, [XBLOCK, RBLOCK]),
        tmp5, other=0.0)
    tmp226 = tmp218 - tmp225
    tmp227 = tl.load(in_ptr23 + tl.broadcast_to(-1 + r0, [XBLOCK, RBLOCK]),
        tmp5, other=0.0)
    tmp228 = tmp218 + tmp227
    tmp229 = tmp226 / tmp228
    tmp230 = tmp13 - tmp229
    tmp231 = tmp224 - tmp230
    tmp232 = tl.full(tmp231.shape, 0.0, tmp231.dtype)
    tmp233 = tl.where(tmp5, tmp231, tmp232)
    tmp235 = tmp218 - tmp234
    tmp237 = tmp218 + tmp236
    tmp238 = tmp235 / tmp237
    tmp239 = tmp13 - tmp238
    tmp240 = tl.where(tmp5, tmp233, tmp239)
    tmp241 = tmp216 * tmp240
    tmp242 = tl.broadcast_to(tmp241, [XBLOCK, RBLOCK])
    tmp244 = tl.sum(tmp242, 1)[:, None]
    tmp245 = tmp34 + tmp64
    tmp246 = tmp245 + tmp94
    tmp247 = tmp246 + tmp124
    tmp248 = 0.25
    tmp249 = tmp247 * tmp248
    tmp250 = tmp154 + tmp184
    tmp251 = tmp250 + tmp214
    tmp252 = tmp251 + tmp244
    tmp253 = tmp252 * tmp248
    tmp254 = tmp249 + tmp253
    tmp255 = 0.5
    tmp256 = tmp254 * tmp255
    tl.debug_barrier()
    tl.store(in_out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp256, None)


def call(args):
    arg0_1, arg1_1 = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(arg1_1, (4, 4, 4, 4), (64, 16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((64,), (1,), torch.float32)
        buf20 = empty_strided_cuda((64,), (1,), torch.float32)
        buf2 = empty_strided_cuda((), (), torch.float32)
        buf3 = empty_strided_cuda((64,), (1,), torch.float32)
        buf4 = empty_strided_cuda((64,), (1,), torch.float32)
        buf22 = empty_strided_cuda((), (), torch.float32)
        buf23 = empty_strided_cuda((64,), (1,), torch.float32)
        buf24 = empty_strided_cuda((64,), (1,), torch.float32)
        get_raw_stream(0)
        triton_per_fused_cumsum_index_mul_rsub_sort_sub_sum_0[grid(1)](arg0_1,
            arg1_1, buf0, buf20, buf2, buf3, buf4, buf22, buf23, buf24, 1, 
            64, XBLOCK=1, num_warps=2, num_stages=1)
        buf5 = empty_strided_cuda((64,), (1,), torch.float32)
        buf25 = empty_strided_cuda((64,), (1,), torch.float32)
        buf7 = empty_strided_cuda((), (), torch.float32)
        buf8 = empty_strided_cuda((64,), (1,), torch.float32)
        buf9 = empty_strided_cuda((64,), (1,), torch.float32)
        buf27 = empty_strided_cuda((), (), torch.float32)
        buf28 = empty_strided_cuda((64,), (1,), torch.float32)
        buf29 = empty_strided_cuda((64,), (1,), torch.float32)
        triton_per_fused_cumsum_index_mul_rsub_sort_sub_sum_1[grid(1)](arg0_1,
            arg1_1, buf5, buf25, buf7, buf8, buf9, buf27, buf28, buf29, 1, 
            64, XBLOCK=1, num_warps=2, num_stages=1)
        buf10 = empty_strided_cuda((64,), (1,), torch.float32)
        buf30 = empty_strided_cuda((64,), (1,), torch.float32)
        buf12 = empty_strided_cuda((), (), torch.float32)
        buf13 = empty_strided_cuda((64,), (1,), torch.float32)
        buf14 = empty_strided_cuda((64,), (1,), torch.float32)
        buf32 = empty_strided_cuda((), (), torch.float32)
        buf33 = empty_strided_cuda((64,), (1,), torch.float32)
        buf34 = empty_strided_cuda((64,), (1,), torch.float32)
        triton_per_fused_cumsum_index_mul_rsub_sort_sub_sum_2[grid(1)](arg0_1,
            arg1_1, buf10, buf30, buf12, buf13, buf14, buf32, buf33, buf34,
            1, 64, XBLOCK=1, num_warps=2, num_stages=1)
        buf15 = empty_strided_cuda((64,), (1,), torch.float32)
        buf35 = empty_strided_cuda((64,), (1,), torch.float32)
        buf17 = empty_strided_cuda((), (), torch.float32)
        buf18 = empty_strided_cuda((64,), (1,), torch.float32)
        buf19 = empty_strided_cuda((64,), (1,), torch.float32)
        buf37 = empty_strided_cuda((), (), torch.float32)
        buf38 = empty_strided_cuda((64,), (1,), torch.float32)
        buf39 = empty_strided_cuda((64,), (1,), torch.float32)
        triton_per_fused_cumsum_index_mul_rsub_sort_sub_sum_3[grid(1)](arg0_1,
            arg1_1, buf15, buf35, buf17, buf18, buf19, buf37, buf38, buf39,
            1, 64, XBLOCK=1, num_warps=2, num_stages=1)
        del arg0_1
        del arg1_1
        buf40 = buf2
        del buf2
        buf41 = buf7
        del buf7
        buf42 = buf12
        del buf12
        buf43 = buf17
        del buf17
        buf44 = buf22
        del buf22
        buf45 = buf27
        del buf27
        buf46 = buf32
        del buf32
        buf47 = buf37
        del buf37
        buf48 = buf40
        del buf40
        triton_per_fused_add_div_dot_relu_rsub_sub_4[grid(1)](buf48, buf41,
            buf42, buf43, buf44, buf45, buf46, buf47, buf0, buf3, buf4,
            buf5, buf8, buf9, buf10, buf13, buf14, buf15, buf18, buf19,
            buf20, buf23, buf24, buf25, buf28, buf29, buf30, buf33, buf34,
            buf35, buf38, buf39, 1, 64, XBLOCK=1, num_warps=2, num_stages=1)
        del buf0
        del buf10
        del buf13
        del buf14
        del buf15
        del buf18
        del buf19
        del buf20
        del buf23
        del buf24
        del buf25
        del buf28
        del buf29
        del buf3
        del buf30
        del buf33
        del buf34
        del buf35
        del buf38
        del buf39
        del buf4
        del buf41
        del buf42
        del buf43
        del buf44
        del buf45
        del buf46
        del buf47
        del buf5
        del buf8
        del buf9
    return buf48,


def jaccard(outputs, targets, per_image=False, non_empty=False, min_pixels=5):
    batch_size = outputs.size()[0]
    eps = 0.001
    if not per_image:
        batch_size = 1
    dice_target = targets.contiguous().view(batch_size, -1).float()
    dice_output = outputs.contiguous().view(batch_size, -1)
    target_sum = torch.sum(dice_target, dim=1)
    intersection = torch.sum(dice_output * dice_target, dim=1)
    losses = 1.0 - (intersection + eps) / (torch.sum(dice_output +
        dice_target, dim=1) - intersection + eps)
    if non_empty:
        assert per_image
        non_empty_images = 0
        sum_loss = 0
        for i in range(batch_size):
            if target_sum[i] > min_pixels:
                sum_loss += losses[i]
                non_empty_images += 1
        if non_empty_images == 0:
            return 0
        else:
            return sum_loss / non_empty_images
    return losses.mean()


def flatten_binary_scores(scores, labels, ignore=None):
    """
    Flattens predictions in the batch (binary case)
    Remove labels equal to 'ignore'
    """
    scores = scores.view(-1)
    labels = labels.view(-1)
    if ignore is None:
        return scores, labels
    valid = labels != ignore
    vscores = scores[valid]
    vlabels = labels[valid]
    return vscores, vlabels


def lovasz_grad(gt_sorted):
    """
    Computes gradient of the Lovasz extension w.r.t sorted errors
    See Alg. 1 in paper
    """
    p = len(gt_sorted)
    gts = gt_sorted.sum()
    intersection = gts.float() - gt_sorted.float().cumsum(0)
    union = gts.float() + (1 - gt_sorted).float().cumsum(0)
    jaccard = 1.0 - intersection / union
    if p > 1:
        jaccard[1:p] = jaccard[1:p] - jaccard[0:-1]
    return jaccard


def lovasz_hinge_flat(logits, labels):
    if len(labels) == 0:
        return logits.sum() * 0.0
    signs = 2.0 * labels.float() - 1.0
    errors = 1.0 - logits * signs
    errors_sorted, perm = torch.sort(errors, dim=0, descending=True)
    perm = perm.data
    gt_sorted = labels[perm]
    grad = lovasz_grad(gt_sorted)
    loss = torch.dot(F.relu(errors_sorted), grad)
    return loss


def mean(l, ignore_nan=False, empty=0):
    """
    nanmean compatible with generators.
    """
    ll = iter(l)
    if ignore_nan:
        ll = filterfalse(np.isnan, ll)
    try:
        n = 1
        acc = next(ll)
    except StopIteration:
        if empty == 'raise':
            raise ValueError('Empty mean')
        return empty
    for n, v in enumerate(ll, 2):
        acc += v
    if n == 1:
        return acc
    return acc / n


def lovasz_hinge(logits, labels, per_image=True, ignore=None):
    if per_image:
        loss = mean(lovasz_hinge_flat(*flatten_binary_scores(log.unsqueeze(
            0), lab.unsqueeze(0), ignore)) for log, lab in zip(logits, labels))
    else:
        loss = lovasz_hinge_flat(*flatten_binary_scores(logits, labels, ignore)
            )
    return loss


def symmetric_lovasz(outputs, targets):
    return (lovasz_hinge(outputs, targets) + lovasz_hinge(-outputs, 1 -
        targets)) / 2


class LovaszLossNew(nn.Module):

    def __init__(self, ignore_index=255, per_image=True):
        super().__init__()
        self.ignore_index = ignore_index
        self.per_image = per_image

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


Measured latency_ms: 0.136751