# 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_4/inductor_cache/d5/cd52mlrkv7fg7ptxw5vaz664ocriv2bgu4unk4oe24v6nbgmv7uv.py
# Topologically Sorted Source Nodes: [mul, intersection, ground_o, pred_o], Original ATen: [aten.mul, aten.sum]
# Source node to ATen node mapping:
#   ground_o => sum_2
#   intersection => sum_1
#   mul => mul
#   pred_o => sum_3
# Graph fragment:
#   %mul : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%arg1_1, %arg0_1), kwargs = {})
#   %sum_1 : [num_users=1] = call_function[target=torch.ops.aten.sum.dim_IntList](args = (%mul, [2, 3]), kwargs = {})
#   %sum_2 : [num_users=2] = call_function[target=torch.ops.aten.sum.dim_IntList](args = (%arg1_1, [2, 3]), kwargs = {})
#   %sum_3 : [num_users=1] = call_function[target=torch.ops.aten.sum.dim_IntList](args = (%arg0_1, [2, 3]), kwargs = {})
triton_per_fused_mul_sum_0 = async_compile.triton('triton_per_fused_mul_sum_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.persistent_reduction(
    size_hints=[16, 16],
    reduction_hint=ReductionHint.INNER,
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: '*fp32', 3: '*fp32', 4: '*fp32', 5: 'i32', 6: '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, 5, 6), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_per_fused_mul_sum_0', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 2, 'num_reduction': 3, '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}
)
@triton.jit
def triton_per_fused_mul_sum_0(in_ptr0, in_ptr1, out_ptr0, out_ptr1, out_ptr2, xnumel, rnumel, XBLOCK : tl.constexpr):
    xnumel = 16
    rnumel = 16
    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, :]
    roffset = 0
    rmask = tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r1 = rindex
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (r1 + (16*x0)), xmask, other=0.0)
    tmp1 = tl.load(in_ptr1 + (r1 + (16*x0)), xmask, other=0.0)
    tmp2 = tmp0 * tmp1
    tmp3 = tl.broadcast_to(tmp2, [XBLOCK, RBLOCK])
    tmp5 = tl.where(xmask, tmp3, 0)
    tmp6 = tl.sum(tmp5, 1)[:, None]
    tmp7 = tl.broadcast_to(tmp0, [XBLOCK, RBLOCK])
    tmp9 = tl.where(xmask, tmp7, 0)
    tmp10 = tl.sum(tmp9, 1)[:, None]
    tmp11 = tl.broadcast_to(tmp1, [XBLOCK, RBLOCK])
    tmp13 = tl.where(xmask, tmp11, 0)
    tmp14 = tl.sum(tmp13, 1)[:, None]
    tl.store(out_ptr0 + (x0), tmp6, xmask)
    tl.store(out_ptr1 + (x0), tmp10, xmask)
    tl.store(out_ptr2 + (x0), tmp14, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_4/inductor_cache/6l/c6l2gkrph5u6qhzfpwcqqfu277drwwd7ra3dyhys4nf64xwbbpre.py
# Topologically Sorted Source Nodes: [setitem], Original ATen: [aten.lift_fresh, aten.index_put]
# Source node to ATen node mapping:
#   setitem => full_default, index_put
# Graph fragment:
#   %full_default : [num_users=1] = call_function[target=torch.ops.aten.full.default](args = ([], 0.0), kwargs = {dtype: torch.float32, layout: torch.strided, device: cpu, pin_memory: False})
#   %index_put : [num_users=1] = call_function[target=torch.ops.aten.index_put.default](args = (%select, [%isinf], %full_default), kwargs = {})
triton_poi_fused_index_put_lift_fresh_1 = async_compile.triton('triton_poi_fused_index_put_lift_fresh_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=[4], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 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), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_index_put_lift_fresh_1', 'mutated_arg_names': [], '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_1(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 4
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (x0), xmask)
    tmp1 = tmp0 * tmp0
    tmp2 = tl.full([1], 1, tl.int32)
    tmp3 = tmp2 / tmp1
    tmp4 = libdevice.isinf(tmp3).to(tl.int1)
    tmp5 = 0.0
    tmp6 = tl.where(tmp4, tmp5, tmp3)
    tl.store(out_ptr0 + (x0), tmp6, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_4/inductor_cache/cy/ccy7idfxu272yvofjbbp54inr6a5zn7jkvgbq7732tfnhhq2lmr7.py
# Topologically Sorted Source Nodes: [mul_1, w], Original ATen: [aten.mul, aten.reciprocal]
# Source node to ATen node mapping:
#   mul_1 => mul_1
#   w => reciprocal
# Graph fragment:
#   %mul_1 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%sum_2, %sum_2), kwargs = {})
#   %reciprocal : [num_users=2] = call_function[target=torch.ops.aten.reciprocal.default](args = (%mul_1,), kwargs = {})
#   %select_scatter_default : [num_users=2] = call_function[target=torch.ops.aten.select_scatter.default](args = (%reciprocal, %index_put, 0, 0), kwargs = {})
triton_poi_fused_mul_reciprocal_2 = async_compile.triton('triton_poi_fused_mul_reciprocal_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=[16], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: '*fp32', 3: '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), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_mul_reciprocal_2', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 2, '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_mul_reciprocal_2(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = (xindex // 4)
    x0 = xindex % 4
    x2 = xindex
    tmp3 = tl.load(in_ptr0 + (x0), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr1 + (x2), xmask)
    tmp0 = x1
    tmp1 = tl.full([1], 0, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp5 = tmp4 * tmp4
    tmp6 = tl.full([1], 1, tl.int32)
    tmp7 = tmp6 / tmp5
    tmp8 = tl.where(tmp2, tmp3, tmp7)
    tl.store(out_ptr0 + (x2), tmp8, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_4/inductor_cache/jk/cjkuxs5m7ejvswg7hrk2llstygqftgn2i6wdk7h25jxqaj665z6b.py
# Topologically Sorted Source Nodes: [max_1, setitem_1], Original ATen: [aten.max, aten.index_put]
# Source node to ATen node mapping:
#   max_1 => max_1
#   setitem_1 => index_put_1
# Graph fragment:
#   %max_1 : [num_users=1] = call_function[target=torch.ops.aten.max.default](args = (%select_4,), kwargs = {})
#   %index_put_1 : [num_users=1] = call_function[target=torch.ops.aten.index_put_.default](args = (%select_4, [%isinf], %max_1), kwargs = {})
triton_per_fused_index_put_max_3 = async_compile.triton('triton_per_fused_index_put_max_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.persistent_reduction(
    size_hints=[1, 4],
    reduction_hint=ReductionHint.INNER,
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: '*fp32', 3: 'i32', 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': {3: 1}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2), equal_to_1=(3,))]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_per_fused_index_put_max_3', 'mutated_arg_names': ['out_ptr2'], 'no_x_dim': False, 'num_load': 2, 'num_reduction': 1, '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}
)
@triton.jit
def triton_per_fused_index_put_max_3(in_ptr0, in_ptr1, out_ptr2, xnumel, rnumel, XBLOCK : tl.constexpr):
    xnumel = 1
    rnumel = 4
    RBLOCK: tl.constexpr = 4
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[None, :]
    roffset = 0
    rmask = tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r0 = rindex
    tmp2 = tl.load(in_ptr0 + (r0), None)
    tmp3 = tl.load(in_ptr1 + (r0), None)
    tmp0 = tl.full([1, 1], 0, tl.int32)
    tmp1 = tmp0 == tmp0
    tmp4 = tmp3 * tmp3
    tmp5 = tl.full([1, 1], 1, tl.int32)
    tmp6 = tmp5 / tmp4
    tmp7 = tl.where(tmp1, tmp2, tmp6)
    tmp8 = tl.broadcast_to(tmp7, [XBLOCK, RBLOCK])
    tmp10 = triton_helpers.max2(tmp8, 1)[:, None]
    tmp11 = libdevice.isinf(tmp6).to(tl.int1)
    tmp12 = tl.where(tmp11, tmp10, tmp7)
    tl.store(out_ptr2 + (tl.broadcast_to(r0, [XBLOCK, RBLOCK])), tmp12, None)
''', device_str='cuda')


# kernel path: runs/run_shard_4/inductor_cache/4q/c4ql4xyg2xlaa3uezk7gjgywmt3fw3tynazq4ckpuwcsgzw6zhei.py
# Topologically Sorted Source Nodes: [], Original ATen: []
# Source node to ATen node mapping:
# Graph fragment:
#   %select_scatter_default_1 : [num_users=2] = call_function[target=torch.ops.aten.select_scatter.default](args = (%select_scatter_default, %index_put_1, 0, 0), kwargs = {})
triton_poi_fused_4 = async_compile.triton('triton_poi_fused_4', '''
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=[16], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 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_4', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 2, '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_4(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = (xindex // 4)
    x0 = xindex % 4
    x2 = xindex
    tmp3 = tl.load(in_ptr0 + (x0), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + (x2), xmask)
    tmp0 = x1
    tmp1 = tl.full([1], 0, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp5 = tl.where(tmp2, tmp3, tmp4)
    tl.store(out_ptr0 + (x2), tmp5, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_4/inductor_cache/o5/co5ocweef7wvpqwak7rmrihdxwrqz6mechmdzkcte5vjlsv5tdhx.py
# Topologically Sorted Source Nodes: [setitem_2], Original ATen: [aten.lift_fresh, aten.index_put]
# Source node to ATen node mapping:
#   setitem_2 => full_default_1, index_put_2
# Graph fragment:
#   %full_default_1 : [num_users=1] = call_function[target=torch.ops.aten.full.default](args = ([], 0.0), kwargs = {dtype: torch.float32, layout: torch.strided, device: cpu, pin_memory: False})
#   %index_put_2 : [num_users=1] = call_function[target=torch.ops.aten.index_put_.default](args = (%select_6, [%isinf_1], %full_default_1), kwargs = {})
triton_poi_fused_index_put_lift_fresh_5 = async_compile.triton('triton_poi_fused_index_put_lift_fresh_5', '''
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=[4], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 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), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_index_put_lift_fresh_5', 'mutated_arg_names': ['out_ptr1'], 'no_x_dim': False, 'num_load': 2, '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_5(in_ptr0, out_ptr1, xnumel, XBLOCK : tl.constexpr):
    xnumel = 4
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp3 = tl.load(in_ptr0 + (x0), xmask)
    tmp4 = tl.load(in_ptr0 + (4 + x0), xmask)
    tmp0 = tl.full([1], 1, tl.int32)
    tmp1 = tl.full([1], 0, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp5 = tl.where(tmp2, tmp3, tmp4)
    tmp6 = libdevice.isinf(tmp5).to(tl.int1)
    tmp7 = 0.0
    tmp8 = tl.where(tmp6, tmp7, tmp5)
    tl.store(out_ptr1 + (4 + x0), tmp8, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_4/inductor_cache/yv/cyvyap3cqstn6kbri4i6qtdbvfru2ywjrkts4h2byghwjatyop4k.py
# Topologically Sorted Source Nodes: [], Original ATen: []
# Source node to ATen node mapping:
# Graph fragment:
#   %select_scatter_default_2 : [num_users=2] = call_function[target=torch.ops.aten.select_scatter.default](args = (%select_scatter_default_1, %index_put_2, 0, 1), kwargs = {})
triton_poi_fused_6 = async_compile.triton('triton_poi_fused_6', '''
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=[16], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 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_6', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 2, '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_6(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = (xindex // 4)
    x0 = xindex % 4
    x2 = xindex
    tmp3 = tl.load(in_ptr0 + (4 + x0), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + (x2), xmask)
    tmp0 = x1
    tmp1 = tl.full([1], 1, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp5 = tl.where(tmp2, tmp3, tmp4)
    tl.store(out_ptr0 + (x2), tmp5, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_4/inductor_cache/tg/ctgfsxxk2urovszhep4mbnqcv7roo2a5o633p6v6krrzhegbz6tr.py
# Topologically Sorted Source Nodes: [max_2, setitem_3], Original ATen: [aten.max, aten.index_put]
# Source node to ATen node mapping:
#   max_2 => max_2
#   setitem_3 => index_put_3
# Graph fragment:
#   %max_2 : [num_users=1] = call_function[target=torch.ops.aten.max.default](args = (%select_7,), kwargs = {})
#   %index_put_3 : [num_users=1] = call_function[target=torch.ops.aten.index_put_.default](args = (%select_7, [%isinf_1], %max_2), kwargs = {})
triton_per_fused_index_put_max_7 = async_compile.triton('triton_per_fused_index_put_max_7', '''
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.persistent_reduction(
    size_hints=[1, 4],
    reduction_hint=ReductionHint.INNER,
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: '*fp32', 3: 'i32', 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': {3: 1}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2), equal_to_1=(3,))]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_per_fused_index_put_max_7', 'mutated_arg_names': ['out_ptr2'], 'no_x_dim': False, 'num_load': 3, 'num_reduction': 1, '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}
)
@triton.jit
def triton_per_fused_index_put_max_7(in_ptr0, in_ptr1, out_ptr2, xnumel, rnumel, XBLOCK : tl.constexpr):
    xnumel = 1
    rnumel = 4
    RBLOCK: tl.constexpr = 4
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[None, :]
    roffset = 0
    rmask = tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r0 = rindex
    tmp2 = tl.load(in_ptr0 + (4 + r0), None)
    tmp9 = tl.load(in_ptr1 + (r0), None)
    tmp10 = tl.load(in_ptr1 + (4 + r0), None)
    tmp0 = tl.full([1, 1], 1, tl.int32)
    tmp1 = tmp0 == tmp0
    tmp3 = tl.where(tmp1, tmp2, tmp2)
    tmp4 = tl.broadcast_to(tmp3, [XBLOCK, RBLOCK])
    tmp6 = triton_helpers.max2(tmp4, 1)[:, None]
    tmp7 = tl.full([1, 1], 0, tl.int32)
    tmp8 = tmp0 == tmp7
    tmp11 = tl.where(tmp8, tmp9, tmp10)
    tmp12 = libdevice.isinf(tmp11).to(tl.int1)
    tmp13 = tl.where(tmp12, tmp6, tmp3)
    tl.store(out_ptr2 + (tl.broadcast_to(4 + r0, [XBLOCK, RBLOCK])), tmp13, None)
''', device_str='cuda')


# kernel path: runs/run_shard_4/inductor_cache/fy/cfy4vgr7piqugeft2crnhfbw5ielxsgean22zjdvpklq4txeer5o.py
# Topologically Sorted Source Nodes: [setitem_4], Original ATen: [aten.lift_fresh, aten.index_put]
# Source node to ATen node mapping:
#   setitem_4 => full_default_2, index_put_4
# Graph fragment:
#   %full_default_2 : [num_users=1] = call_function[target=torch.ops.aten.full.default](args = ([], 0.0), kwargs = {dtype: torch.float32, layout: torch.strided, device: cpu, pin_memory: False})
#   %index_put_4 : [num_users=1] = call_function[target=torch.ops.aten.index_put_.default](args = (%select_9, [%isinf_2], %full_default_2), kwargs = {})
triton_poi_fused_index_put_lift_fresh_8 = async_compile.triton('triton_poi_fused_index_put_lift_fresh_8', '''
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=[4], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 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), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_index_put_lift_fresh_8', 'mutated_arg_names': ['out_ptr1'], 'no_x_dim': False, 'num_load': 2, '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_8(in_ptr0, out_ptr1, xnumel, XBLOCK : tl.constexpr):
    xnumel = 4
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp3 = tl.load(in_ptr0 + (4 + x0), xmask)
    tmp4 = tl.load(in_ptr0 + (8 + x0), xmask)
    tmp0 = tl.full([1], 2, tl.int32)
    tmp1 = tl.full([1], 1, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp5 = tl.where(tmp2, tmp3, tmp4)
    tmp6 = libdevice.isinf(tmp5).to(tl.int1)
    tmp7 = 0.0
    tmp8 = tl.where(tmp6, tmp7, tmp5)
    tl.store(out_ptr1 + (8 + x0), tmp8, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_4/inductor_cache/zj/czjcmriyujs5brmyh7stjrkfk37e4twzbljor3wofssbrxsuw63t.py
# Topologically Sorted Source Nodes: [], Original ATen: []
# Source node to ATen node mapping:
# Graph fragment:
#   %select_scatter_default_4 : [num_users=2] = call_function[target=torch.ops.aten.select_scatter.default](args = (%select_scatter_default_3, %index_put_4, 0, 2), kwargs = {})
triton_poi_fused_9 = async_compile.triton('triton_poi_fused_9', '''
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=[16], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 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_9', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 2, '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_9(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = (xindex // 4)
    x0 = xindex % 4
    x2 = xindex
    tmp3 = tl.load(in_ptr0 + (8 + x0), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + (x2), xmask)
    tmp0 = x1
    tmp1 = tl.full([1], 2, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp5 = tl.where(tmp2, tmp3, tmp4)
    tl.store(out_ptr0 + (x2), tmp5, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_4/inductor_cache/hx/chxbr5aueebc6hjsn4ej4rftudtvp5pbxyu5cbojq5zy2lksc2pa.py
# Topologically Sorted Source Nodes: [max_3, setitem_5], Original ATen: [aten.max, aten.index_put]
# Source node to ATen node mapping:
#   max_3 => max_3
#   setitem_5 => index_put_5
# Graph fragment:
#   %max_3 : [num_users=1] = call_function[target=torch.ops.aten.max.default](args = (%select_10,), kwargs = {})
#   %index_put_5 : [num_users=1] = call_function[target=torch.ops.aten.index_put_.default](args = (%select_10, [%isinf_2], %max_3), kwargs = {})
triton_per_fused_index_put_max_10 = async_compile.triton('triton_per_fused_index_put_max_10', '''
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.persistent_reduction(
    size_hints=[1, 4],
    reduction_hint=ReductionHint.INNER,
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: '*fp32', 3: 'i32', 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': {3: 1}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2), equal_to_1=(3,))]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_per_fused_index_put_max_10', 'mutated_arg_names': ['out_ptr2'], 'no_x_dim': False, 'num_load': 3, 'num_reduction': 1, '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}
)
@triton.jit
def triton_per_fused_index_put_max_10(in_ptr0, in_ptr1, out_ptr2, xnumel, rnumel, XBLOCK : tl.constexpr):
    xnumel = 1
    rnumel = 4
    RBLOCK: tl.constexpr = 4
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[None, :]
    roffset = 0
    rmask = tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r0 = rindex
    tmp2 = tl.load(in_ptr0 + (8 + r0), None)
    tmp9 = tl.load(in_ptr1 + (4 + r0), None)
    tmp10 = tl.load(in_ptr1 + (8 + r0), None)
    tmp0 = tl.full([1, 1], 2, tl.int32)
    tmp1 = tmp0 == tmp0
    tmp3 = tl.where(tmp1, tmp2, tmp2)
    tmp4 = tl.broadcast_to(tmp3, [XBLOCK, RBLOCK])
    tmp6 = triton_helpers.max2(tmp4, 1)[:, None]
    tmp7 = tl.full([1, 1], 1, tl.int32)
    tmp8 = tmp0 == tmp7
    tmp11 = tl.where(tmp8, tmp9, tmp10)
    tmp12 = libdevice.isinf(tmp11).to(tl.int1)
    tmp13 = tl.where(tmp12, tmp6, tmp3)
    tl.store(out_ptr2 + (tl.broadcast_to(8 + r0, [XBLOCK, RBLOCK])), tmp13, None)
''', device_str='cuda')


# kernel path: runs/run_shard_4/inductor_cache/wq/cwqjjijigoiri7ebhb452bt4tizigzg2o6ji6lsekghkviga7xp3.py
# Topologically Sorted Source Nodes: [setitem_6], Original ATen: [aten.lift_fresh, aten.index_put]
# Source node to ATen node mapping:
#   setitem_6 => full_default_3, index_put_6
# Graph fragment:
#   %full_default_3 : [num_users=1] = call_function[target=torch.ops.aten.full.default](args = ([], 0.0), kwargs = {dtype: torch.float32, layout: torch.strided, device: cpu, pin_memory: False})
#   %index_put_6 : [num_users=1] = call_function[target=torch.ops.aten.index_put_.default](args = (%select_12, [%isinf_3], %full_default_3), kwargs = {})
triton_poi_fused_index_put_lift_fresh_11 = async_compile.triton('triton_poi_fused_index_put_lift_fresh_11', '''
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=[4], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 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), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_index_put_lift_fresh_11', 'mutated_arg_names': ['out_ptr1'], 'no_x_dim': False, 'num_load': 2, '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_11(in_ptr0, out_ptr1, xnumel, XBLOCK : tl.constexpr):
    xnumel = 4
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp3 = tl.load(in_ptr0 + (8 + x0), xmask)
    tmp4 = tl.load(in_ptr0 + (12 + x0), xmask)
    tmp0 = tl.full([1], 3, tl.int32)
    tmp1 = tl.full([1], 2, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp5 = tl.where(tmp2, tmp3, tmp4)
    tmp6 = libdevice.isinf(tmp5).to(tl.int1)
    tmp7 = 0.0
    tmp8 = tl.where(tmp6, tmp7, tmp5)
    tl.store(out_ptr1 + (12 + x0), tmp8, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_4/inductor_cache/ux/cuxhakfkc47pmz4gov3quzmpxdbvtixikhnf3nicsxfr6rvsdnli.py
# Topologically Sorted Source Nodes: [], Original ATen: []
# Source node to ATen node mapping:
# Graph fragment:
#   %select_scatter_default_6 : [num_users=2] = call_function[target=torch.ops.aten.select_scatter.default](args = (%select_scatter_default_5, %index_put_6, 0, 3), kwargs = {})
triton_poi_fused_12 = async_compile.triton('triton_poi_fused_12', '''
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=[16], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 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_12', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 2, '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_12(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = (xindex // 4)
    x0 = xindex % 4
    x2 = xindex
    tmp3 = tl.load(in_ptr0 + (12 + x0), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + (x2), xmask)
    tmp0 = x1
    tmp1 = tl.full([1], 3, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp5 = tl.where(tmp2, tmp3, tmp4)
    tl.store(out_ptr0 + (x2), tmp5, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_4/inductor_cache/lf/clf4pxfzowgmhow2bm4nhu2drefrnqqu6uqv5egtoum633aeciwm.py
# Topologically Sorted Source Nodes: [max_4, setitem_7], Original ATen: [aten.max, aten.index_put]
# Source node to ATen node mapping:
#   max_4 => max_4
#   setitem_7 => index_put_7
# Graph fragment:
#   %max_4 : [num_users=1] = call_function[target=torch.ops.aten.max.default](args = (%select_13,), kwargs = {})
#   %index_put_7 : [num_users=1] = call_function[target=torch.ops.aten.index_put_.default](args = (%select_13, [%isinf_3], %max_4), kwargs = {})
triton_per_fused_index_put_max_13 = async_compile.triton('triton_per_fused_index_put_max_13', '''
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.persistent_reduction(
    size_hints=[1, 4],
    reduction_hint=ReductionHint.INNER,
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: '*fp32', 3: 'i32', 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': {3: 1}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2), equal_to_1=(3,))]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_per_fused_index_put_max_13', 'mutated_arg_names': ['out_ptr2'], 'no_x_dim': False, 'num_load': 3, 'num_reduction': 1, '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}
)
@triton.jit
def triton_per_fused_index_put_max_13(in_ptr0, in_ptr1, out_ptr2, xnumel, rnumel, XBLOCK : tl.constexpr):
    xnumel = 1
    rnumel = 4
    RBLOCK: tl.constexpr = 4
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[None, :]
    roffset = 0
    rmask = tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r0 = rindex
    tmp2 = tl.load(in_ptr0 + (12 + r0), None)
    tmp9 = tl.load(in_ptr1 + (8 + r0), None)
    tmp10 = tl.load(in_ptr1 + (12 + r0), None)
    tmp0 = tl.full([1, 1], 3, tl.int32)
    tmp1 = tmp0 == tmp0
    tmp3 = tl.where(tmp1, tmp2, tmp2)
    tmp4 = tl.broadcast_to(tmp3, [XBLOCK, RBLOCK])
    tmp6 = triton_helpers.max2(tmp4, 1)[:, None]
    tmp7 = tl.full([1, 1], 2, tl.int32)
    tmp8 = tmp0 == tmp7
    tmp11 = tl.where(tmp8, tmp9, tmp10)
    tmp12 = libdevice.isinf(tmp11).to(tl.int1)
    tmp13 = tl.where(tmp12, tmp6, tmp3)
    tl.store(out_ptr2 + (tl.broadcast_to(12 + r0, [XBLOCK, RBLOCK])), tmp13, None)
''', device_str='cuda')


# kernel path: runs/run_shard_4/inductor_cache/s2/cs2wyfvtdb7vmxx7u5otf7upxsggzrwonnusuhl3h5txytfu73v3.py
# Topologically Sorted Source Nodes: [mul_2, sum_4, mul_3, add_1, denominator, mul_4, sum_5, add_2, truediv, f, f_1], Original ATen: [aten.mul, aten.sum, aten.add, aten.div, aten.rsub, aten.mean]
# Source node to ATen node mapping:
#   add_1 => add_1
#   add_2 => add_2
#   denominator => add
#   f => sub
#   f_1 => mean
#   mul_2 => mul_2
#   mul_3 => mul_3
#   mul_4 => mul_4
#   sum_4 => sum_4
#   sum_5 => sum_5
#   truediv => div
# Graph fragment:
#   %select_scatter_default_7 : [num_users=2] = call_function[target=torch.ops.aten.select_scatter.default](args = (%select_scatter_default_6, %index_put_7, 0, 3), kwargs = {})
#   %mul_2 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%sum_1, %select_scatter_default_7), kwargs = {})
#   %sum_4 : [num_users=1] = call_function[target=torch.ops.aten.sum.dim_IntList](args = (%mul_2, [1]), kwargs = {})
#   %mul_3 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%sum_4, 2.0), kwargs = {})
#   %add_1 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul_3, 1e-05), kwargs = {})
#   %add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%sum_2, %sum_3), kwargs = {})
#   %mul_4 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%add, %select_scatter_default_7), kwargs = {})
#   %sum_5 : [num_users=1] = call_function[target=torch.ops.aten.sum.dim_IntList](args = (%mul_4, [1]), kwargs = {})
#   %add_2 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%sum_5, 1e-05), kwargs = {})
#   %div : [num_users=1] = call_function[target=torch.ops.aten.div.Tensor](args = (%add_1, %add_2), kwargs = {})
#   %sub : [num_users=1] = call_function[target=torch.ops.aten.sub.Tensor](args = (1.0, %div), kwargs = {})
#   %mean : [num_users=1] = call_function[target=torch.ops.aten.mean.default](args = (%sub,), kwargs = {})
triton_per_fused_add_div_mean_mul_rsub_sum_14 = async_compile.triton('triton_per_fused_add_div_mean_mul_rsub_sum_14', '''
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.persistent_reduction(
    size_hints=[1, 4],
    reduction_hint=ReductionHint.DEFAULT,
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: '*fp32', 3: '*fp32', 4: '*fp32', 5: 'i32', 6: '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': {5: 1}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2, 3, 4), equal_to_1=(5,))]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_per_fused_add_div_mean_mul_rsub_sum_14', 'mutated_arg_names': ['in_out_ptr0'], 'no_x_dim': False, 'num_load': 20, 'num_reduction': 1, '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}
)
@triton.jit
def triton_per_fused_add_div_mean_mul_rsub_sum_14(in_out_ptr0, in_ptr0, in_ptr1, in_ptr2, in_ptr3, xnumel, rnumel, XBLOCK : tl.constexpr):
    xnumel = 1
    rnumel = 4
    RBLOCK: tl.constexpr = 4
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[None, :]
    roffset = 0
    rmask = tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r0 = rindex
    tmp0 = tl.load(in_ptr0 + (4*r0), None, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr1 + (12))
    tmp5 = tl.broadcast_to(tmp4, [XBLOCK, RBLOCK])
    tmp6 = tl.load(in_ptr1 + (4*r0), None, eviction_policy='evict_last')
    tmp9 = tl.load(in_ptr0 + (1 + (4*r0)), None, eviction_policy='evict_last')
    tmp10 = tl.load(in_ptr1 + (13))
    tmp11 = tl.broadcast_to(tmp10, [XBLOCK, RBLOCK])
    tmp12 = tl.load(in_ptr1 + (1 + (4*r0)), None, eviction_policy='evict_last')
    tmp16 = tl.load(in_ptr0 + (2 + (4*r0)), None, eviction_policy='evict_last')
    tmp17 = tl.load(in_ptr1 + (14))
    tmp18 = tl.broadcast_to(tmp17, [XBLOCK, RBLOCK])
    tmp19 = tl.load(in_ptr1 + (2 + (4*r0)), None, eviction_policy='evict_last')
    tmp23 = tl.load(in_ptr0 + (3 + (4*r0)), None, eviction_policy='evict_last')
    tmp24 = tl.load(in_ptr1 + (15))
    tmp25 = tl.broadcast_to(tmp24, [XBLOCK, RBLOCK])
    tmp26 = tl.load(in_ptr1 + (3 + (4*r0)), None, eviction_policy='evict_last')
    tmp30 = tl.load(in_ptr2 + (4*r0), None, eviction_policy='evict_last')
    tmp31 = tl.load(in_ptr3 + (4*r0), None, eviction_policy='evict_last')
    tmp34 = tl.load(in_ptr2 + (1 + (4*r0)), None, eviction_policy='evict_last')
    tmp35 = tl.load(in_ptr3 + (1 + (4*r0)), None, eviction_policy='evict_last')
    tmp39 = tl.load(in_ptr2 + (2 + (4*r0)), None, eviction_policy='evict_last')
    tmp40 = tl.load(in_ptr3 + (2 + (4*r0)), None, eviction_policy='evict_last')
    tmp44 = tl.load(in_ptr2 + (3 + (4*r0)), None, eviction_policy='evict_last')
    tmp45 = tl.load(in_ptr3 + (3 + (4*r0)), None, eviction_policy='evict_last')
    tmp1 = r0
    tmp2 = tl.full([1, 1], 3, tl.int32)
    tmp3 = tmp1 == tmp2
    tmp7 = tl.where(tmp3, tmp5, tmp6)
    tmp8 = tmp0 * tmp7
    tmp13 = tl.where(tmp3, tmp11, tmp12)
    tmp14 = tmp9 * tmp13
    tmp15 = tmp8 + tmp14
    tmp20 = tl.where(tmp3, tmp18, tmp19)
    tmp21 = tmp16 * tmp20
    tmp22 = tmp15 + tmp21
    tmp27 = tl.where(tmp3, tmp25, tmp26)
    tmp28 = tmp23 * tmp27
    tmp29 = tmp22 + tmp28
    tmp32 = tmp30 + tmp31
    tmp33 = tmp32 * tmp7
    tmp36 = tmp34 + tmp35
    tmp37 = tmp36 * tmp13
    tmp38 = tmp33 + tmp37
    tmp41 = tmp39 + tmp40
    tmp42 = tmp41 * tmp20
    tmp43 = tmp38 + tmp42
    tmp46 = tmp44 + tmp45
    tmp47 = tmp46 * tmp27
    tmp48 = tmp43 + tmp47
    tmp49 = 2.0
    tmp50 = tmp29 * tmp49
    tmp51 = 1e-05
    tmp52 = tmp50 + tmp51
    tmp53 = tmp48 + tmp51
    tmp54 = tmp52 / tmp53
    tmp55 = 1.0
    tmp56 = tmp55 - tmp54
    tmp57 = tl.broadcast_to(tmp56, [XBLOCK, RBLOCK])
    tmp59 = tl.sum(tmp57, 1)[:, None]
    tmp60 = 4.0
    tmp61 = tmp59 / tmp60
    tl.debug_barrier()
    tl.store(in_out_ptr0 + (tl.full([XBLOCK, 1], 0, tl.int32)), tmp61, None)
''', device_str='cuda')


async_compile.wait(globals())
del async_compile

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((4, 4), (4, 1), torch.float32)
        buf1 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        buf29 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        # Topologically Sorted Source Nodes: [mul, intersection, ground_o, pred_o], Original ATen: [aten.mul, aten.sum]
        stream0 = get_raw_stream(0)
        triton_per_fused_mul_sum_0.run(arg1_1, arg0_1, buf0, buf1, buf29, 16, 16, grid=grid(16), stream=stream0)
        del arg0_1
        del arg1_1
        buf2 = empty_strided_cuda((4, ), (1, ), torch.float32)
        # Topologically Sorted Source Nodes: [setitem], Original ATen: [aten.lift_fresh, aten.index_put]
        triton_poi_fused_index_put_lift_fresh_1.run(buf1, buf2, 4, grid=grid(4), stream=stream0)
        buf4 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        # Topologically Sorted Source Nodes: [mul_1, w], Original ATen: [aten.mul, aten.reciprocal]
        triton_poi_fused_mul_reciprocal_2.run(buf2, buf1, buf4, 16, grid=grid(16), stream=stream0)
        # Topologically Sorted Source Nodes: [max_1, setitem_1], Original ATen: [aten.max, aten.index_put]
        triton_per_fused_index_put_max_3.run(buf2, buf1, buf4, 1, 4, grid=grid(1), stream=stream0)
        del buf2
        buf7 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        # Topologically Sorted Source Nodes: [], Original ATen: []
        triton_poi_fused_4.run(buf4, buf7, 16, grid=grid(16), stream=stream0)
        # Topologically Sorted Source Nodes: [setitem_2], Original ATen: [aten.lift_fresh, aten.index_put]
        triton_poi_fused_index_put_lift_fresh_5.run(buf4, buf7, 4, grid=grid(4), stream=stream0)
        buf11 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        # Topologically Sorted Source Nodes: [], Original ATen: []
        triton_poi_fused_6.run(buf7, buf11, 16, grid=grid(16), stream=stream0)
        # Topologically Sorted Source Nodes: [max_2, setitem_3], Original ATen: [aten.max, aten.index_put]
        triton_per_fused_index_put_max_7.run(buf7, buf4, buf11, 1, 4, grid=grid(1), stream=stream0)
        buf14 = buf7; del buf7  # reuse
        # Topologically Sorted Source Nodes: [], Original ATen: []
        triton_poi_fused_6.run(buf11, buf14, 16, grid=grid(16), stream=stream0)
        # Topologically Sorted Source Nodes: [setitem_4], Original ATen: [aten.lift_fresh, aten.index_put]
        triton_poi_fused_index_put_lift_fresh_8.run(buf11, buf14, 4, grid=grid(4), stream=stream0)
        buf18 = buf4; del buf4  # reuse
        # Topologically Sorted Source Nodes: [], Original ATen: []
        triton_poi_fused_9.run(buf14, buf18, 16, grid=grid(16), stream=stream0)
        # Topologically Sorted Source Nodes: [max_3, setitem_5], Original ATen: [aten.max, aten.index_put]
        triton_per_fused_index_put_max_10.run(buf14, buf11, buf18, 1, 4, grid=grid(1), stream=stream0)
        buf21 = buf14; del buf14  # reuse
        # Topologically Sorted Source Nodes: [], Original ATen: []
        triton_poi_fused_9.run(buf18, buf21, 16, grid=grid(16), stream=stream0)
        # Topologically Sorted Source Nodes: [setitem_6], Original ATen: [aten.lift_fresh, aten.index_put]
        triton_poi_fused_index_put_lift_fresh_11.run(buf18, buf21, 4, grid=grid(4), stream=stream0)
        buf25 = buf11; del buf11  # reuse
        # Topologically Sorted Source Nodes: [], Original ATen: []
        triton_poi_fused_12.run(buf21, buf25, 16, grid=grid(16), stream=stream0)
        # Topologically Sorted Source Nodes: [max_4, setitem_7], Original ATen: [aten.max, aten.index_put]
        triton_per_fused_index_put_max_13.run(buf21, buf18, buf25, 1, 4, grid=grid(1), stream=stream0)
        del buf18
        del buf21
        buf31 = empty_strided_cuda((), (), torch.float32)
        buf32 = buf31; del buf31  # reuse
        # Topologically Sorted Source Nodes: [mul_2, sum_4, mul_3, add_1, denominator, mul_4, sum_5, add_2, truediv, f, f_1], Original ATen: [aten.mul, aten.sum, aten.add, aten.div, aten.rsub, aten.mean]
        triton_per_fused_add_div_mean_mul_rsub_sum_14.run(buf32, buf0, buf25, buf1, buf29, 1, 4, grid=grid(1), stream=stream0)
        del buf0
        del buf1
        del buf25
        del buf29
    return (buf32, )


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)
    arg1_1 = rand_strided((4, 4, 4, 4), (64, 16, 4, 1), device='cuda:0', dtype=torch.float32)
    fn = lambda: call([arg0_1, arg1_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 collections
import torch
import warnings
from typing import Optional
from typing import Union
from typing import Any
from typing import Callable
from typing import Tuple
import torch.nn
from torch.nn.modules.loss import _Loss
from enum import Enum
import collections.abc


def issequenceiterable(obj: 'Any') ->bool:
    """
    Determine if the object is an iterable sequence and is not a string.
    """
    if torch.is_tensor(obj):
        return int(obj.dim()) > 0
    return isinstance(obj, collections.abc.Iterable) and not isinstance(obj,
        str)


def ensure_tuple(vals: 'Any') ->Tuple[Any, ...]:
    """
    Returns a tuple of `vals`.
    """
    if not issequenceiterable(vals):
        vals = vals,
    return tuple(vals)


def ensure_tuple_size(tup: 'Any', dim: 'int', pad_val: 'Any'=0) ->Tuple[Any,
    ...]:
    """
    Returns a copy of `tup` with `dim` values by either shortened or padded with `pad_val` as necessary.
    """
    tup = ensure_tuple(tup) + (pad_val,) * dim
    return tuple(tup[:dim])


def one_hot(labels: 'torch.Tensor', num_classes: 'int', dtype:
    'torch.dtype'=torch.float, dim: 'int'=1) ->torch.Tensor:
    """
    For a tensor `labels` of dimensions B1[spatial_dims], return a tensor of dimensions `BN[spatial_dims]`
    for `num_classes` N number of classes.

    Example:

        For every value v = labels[b,1,h,w], the value in the result at [b,v,h,w] will be 1 and all others 0.
        Note that this will include the background label, thus a binary mask should be treated as having 2 classes.
    """
    assert labels.dim() > 0, 'labels should have dim of 1 or more.'
    if labels.ndim < dim + 1:
        shape = ensure_tuple_size(labels.shape, dim + 1, 1)
        labels = labels.reshape(*shape)
    sh = list(labels.shape)
    assert sh[dim
        ] == 1, 'labels should have a channel with length equals to one.'
    sh[dim] = num_classes
    o = torch.zeros(size=sh, dtype=dtype, device=labels.device)
    labels = o.scatter_(dim=dim, index=labels.long(), value=1)
    return labels


class LossReduction(Enum):
    """
    See also:
        - :py:class:`monai.losses.dice.DiceLoss`
        - :py:class:`monai.losses.dice.GeneralizedDiceLoss`
        - :py:class:`monai.losses.focal_loss.FocalLoss`
        - :py:class:`monai.losses.tversky.TverskyLoss`
    """
    NONE = 'none'
    MEAN = 'mean'
    SUM = 'sum'


class Weight(Enum):
    """
    See also: :py:class:`monai.losses.dice.GeneralizedDiceLoss`
    """
    SQUARE = 'square'
    SIMPLE = 'simple'
    UNIFORM = 'uniform'


class GeneralizedDiceLoss(_Loss):
    """
    Compute the generalised Dice loss defined in:

        Sudre, C. et. al. (2017) Generalised Dice overlap as a deep learning
        loss function for highly unbalanced segmentations. DLMIA 2017.

    Adapted from:
        https://github.com/NifTK/NiftyNet/blob/v0.6.0/niftynet/layer/loss_segmentation.py#L279
    """

    def __init__(self, include_background: 'bool'=True, to_onehot_y: 'bool'
        =False, sigmoid: 'bool'=False, softmax: 'bool'=False, other_act:
        'Optional[Callable]'=None, w_type: 'Union[Weight, str]'=Weight.
        SQUARE, reduction: 'Union[LossReduction, str]'=LossReduction.MEAN,
        smooth_nr: 'float'=1e-05, smooth_dr: 'float'=1e-05, batch: 'bool'=False
        ) ->None:
        """
        Args:
            include_background: If False channel index 0 (background category) is excluded from the calculation.
            to_onehot_y: whether to convert `y` into the one-hot format. Defaults to False.
            sigmoid: If True, apply a sigmoid function to the prediction.
            softmax: If True, apply a softmax function to the prediction.
            other_act: if don't want to use `sigmoid` or `softmax`, use other callable function to execute
                other activation layers, Defaults to ``None``. for example:
                `other_act = torch.tanh`.
            squared_pred: use squared versions of targets and predictions in the denominator or not.
            w_type: {``"square"``, ``"simple"``, ``"uniform"``}
                Type of function to transform ground truth volume to a weight factor. Defaults to ``"square"``.
            reduction: {``"none"``, ``"mean"``, ``"sum"``}
                Specifies the reduction to apply to the output. Defaults to ``"mean"``.

                - ``"none"``: no reduction will be applied.
                - ``"mean"``: the sum of the output will be divided by the number of elements in the output.
                - ``"sum"``: the output will be summed.
            smooth_nr: a small constant added to the numerator to avoid zero.
            smooth_dr: a small constant added to the denominator to avoid nan.
            batch: whether to sum the intersection and union areas over the batch dimension before the dividing.
                Defaults to False, intersection over union is computed from each item in the batch.

        Raises:
            TypeError: When ``other_act`` is not an ``Optional[Callable]``.
            ValueError: When more than 1 of [``sigmoid=True``, ``softmax=True``, ``other_act is not None``].
                Incompatible values.

        """
        super().__init__(reduction=LossReduction(reduction).value)
        if other_act is not None and not callable(other_act):
            raise TypeError(
                f'other_act must be None or callable but is {type(other_act).__name__}.'
                )
        if int(sigmoid) + int(softmax) + int(other_act is not None) > 1:
            raise ValueError(
                'Incompatible values: more than 1 of [sigmoid=True, softmax=True, other_act is not None].'
                )
        self.include_background = include_background
        self.to_onehot_y = to_onehot_y
        self.sigmoid = sigmoid
        self.softmax = softmax
        self.other_act = other_act
        w_type = Weight(w_type)
        self.w_func: 'Callable' = torch.ones_like
        if w_type == Weight.SIMPLE:
            self.w_func = torch.reciprocal
        elif w_type == Weight.SQUARE:
            self.w_func = lambda x: torch.reciprocal(x * x)
        self.smooth_nr = float(smooth_nr)
        self.smooth_dr = float(smooth_dr)
        self.batch = batch

    def forward(self, input: 'torch.Tensor', target: 'torch.Tensor'
        ) ->torch.Tensor:
        """
        Args:
            input: the shape should be BNH[WD].
            target: the shape should be BNH[WD].

        Raises:
            ValueError: When ``self.reduction`` is not one of ["mean", "sum", "none"].

        """
        if self.sigmoid:
            input = torch.sigmoid(input)
        n_pred_ch = input.shape[1]
        if self.softmax:
            if n_pred_ch == 1:
                warnings.warn(
                    'single channel prediction, `softmax=True` ignored.')
            else:
                input = torch.softmax(input, 1)
        if self.other_act is not None:
            input = self.other_act(input)
        if self.to_onehot_y:
            if n_pred_ch == 1:
                warnings.warn(
                    'single channel prediction, `to_onehot_y=True` ignored.')
            else:
                target = one_hot(target, num_classes=n_pred_ch)
        if not self.include_background:
            if n_pred_ch == 1:
                warnings.warn(
                    'single channel prediction, `include_background=False` ignored.'
                    )
            else:
                target = target[:, 1:]
                input = input[:, 1:]
        assert target.shape == input.shape, f'ground truth has differing shape ({target.shape}) from input ({input.shape})'
        reduce_axis = list(range(2, len(input.shape)))
        if self.batch:
            reduce_axis = [0] + reduce_axis
        intersection = torch.sum(target * input, reduce_axis)
        ground_o = torch.sum(target, reduce_axis)
        pred_o = torch.sum(input, reduce_axis)
        denominator = ground_o + pred_o
        w = self.w_func(ground_o.float())
        for b in w:
            infs = torch.isinf(b)
            b[infs] = 0.0
            b[infs] = torch.max(b)
        f: 'torch.Tensor' = 1.0 - (2.0 * (intersection * w).sum(1) + self.
            smooth_nr) / ((denominator * w).sum(1) + self.smooth_dr)
        if self.reduction == LossReduction.MEAN.value:
            f = torch.mean(f)
        elif self.reduction == LossReduction.SUM.value:
            f = torch.sum(f)
        elif self.reduction == LossReduction.NONE.value:
            pass
        else:
            raise ValueError(
                f'Unsupported reduction: {self.reduction}, available options are ["mean", "sum", "none"].'
                )
        return f


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


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

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
from torch._inductor.runtime.triton_helpers import libdevice
import collections
from typing import Optional
from typing import Union
from typing import Any
from typing import Callable
from typing import Tuple
import torch.nn
from torch.nn.modules.loss import _Loss
from enum import Enum
import collections.abc
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_mul_sum_0(in_ptr0, in_ptr1, out_ptr0, out_ptr1,
    out_ptr2, xnumel, rnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    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 + 16 * x0), xmask, other=0.0)
    tmp1 = tl.load(in_ptr1 + (r1 + 16 * x0), xmask, other=0.0)
    tmp2 = tmp0 * tmp1
    tmp3 = tl.broadcast_to(tmp2, [XBLOCK, RBLOCK])
    tmp5 = tl.where(xmask, tmp3, 0)
    tmp6 = tl.sum(tmp5, 1)[:, None]
    tmp7 = tl.broadcast_to(tmp0, [XBLOCK, RBLOCK])
    tmp9 = tl.where(xmask, tmp7, 0)
    tmp10 = tl.sum(tmp9, 1)[:, None]
    tmp11 = tl.broadcast_to(tmp1, [XBLOCK, RBLOCK])
    tmp13 = tl.where(xmask, tmp11, 0)
    tmp14 = tl.sum(tmp13, 1)[:, None]
    tl.store(out_ptr0 + x0, tmp6, xmask)
    tl.store(out_ptr1 + x0, tmp10, xmask)
    tl.store(out_ptr2 + x0, tmp14, xmask)


@triton.jit
def triton_poi_fused_index_put_lift_fresh_1(in_ptr0, out_ptr0, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 4
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + x0, xmask)
    tmp1 = tmp0 * tmp0
    tmp2 = tl.full([1], 1, tl.int32)
    tmp3 = tmp2 / tmp1
    tmp4 = libdevice.isinf(tmp3).to(tl.int1)
    tmp5 = 0.0
    tmp6 = tl.where(tmp4, tmp5, tmp3)
    tl.store(out_ptr0 + x0, tmp6, xmask)


@triton.jit
def triton_poi_fused_mul_reciprocal_2(in_ptr0, in_ptr1, out_ptr0, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4
    x0 = xindex % 4
    x2 = xindex
    tmp3 = tl.load(in_ptr0 + x0, xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr1 + x2, xmask)
    tmp0 = x1
    tmp1 = tl.full([1], 0, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp5 = tmp4 * tmp4
    tmp6 = tl.full([1], 1, tl.int32)
    tmp7 = tmp6 / tmp5
    tmp8 = tl.where(tmp2, tmp3, tmp7)
    tl.store(out_ptr0 + x2, tmp8, xmask)


@triton.jit
def triton_per_fused_index_put_max_3(in_ptr0, in_ptr1, out_ptr2, xnumel,
    rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 4
    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
    tmp2 = tl.load(in_ptr0 + r0, None)
    tmp3 = tl.load(in_ptr1 + r0, None)
    tmp0 = tl.full([1, 1], 0, tl.int32)
    tmp1 = tmp0 == tmp0
    tmp4 = tmp3 * tmp3
    tmp5 = tl.full([1, 1], 1, tl.int32)
    tmp6 = tmp5 / tmp4
    tmp7 = tl.where(tmp1, tmp2, tmp6)
    tmp8 = tl.broadcast_to(tmp7, [XBLOCK, RBLOCK])
    tmp10 = triton_helpers.max2(tmp8, 1)[:, None]
    tmp11 = libdevice.isinf(tmp6).to(tl.int1)
    tmp12 = tl.where(tmp11, tmp10, tmp7)
    tl.store(out_ptr2 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp12, None)


@triton.jit
def triton_poi_fused_4(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4
    x0 = xindex % 4
    x2 = xindex
    tmp3 = tl.load(in_ptr0 + x0, xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + x2, xmask)
    tmp0 = x1
    tmp1 = tl.full([1], 0, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp5 = tl.where(tmp2, tmp3, tmp4)
    tl.store(out_ptr0 + x2, tmp5, xmask)


@triton.jit
def triton_poi_fused_index_put_lift_fresh_5(in_ptr0, out_ptr1, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 4
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp3 = tl.load(in_ptr0 + x0, xmask)
    tmp4 = tl.load(in_ptr0 + (4 + x0), xmask)
    tmp0 = tl.full([1], 1, tl.int32)
    tmp1 = tl.full([1], 0, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp5 = tl.where(tmp2, tmp3, tmp4)
    tmp6 = libdevice.isinf(tmp5).to(tl.int1)
    tmp7 = 0.0
    tmp8 = tl.where(tmp6, tmp7, tmp5)
    tl.store(out_ptr1 + (4 + x0), tmp8, xmask)


@triton.jit
def triton_poi_fused_6(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4
    x0 = xindex % 4
    x2 = xindex
    tmp3 = tl.load(in_ptr0 + (4 + x0), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + x2, xmask)
    tmp0 = x1
    tmp1 = tl.full([1], 1, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp5 = tl.where(tmp2, tmp3, tmp4)
    tl.store(out_ptr0 + x2, tmp5, xmask)


@triton.jit
def triton_per_fused_index_put_max_7(in_ptr0, in_ptr1, out_ptr2, xnumel,
    rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 4
    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
    tmp2 = tl.load(in_ptr0 + (4 + r0), None)
    tmp9 = tl.load(in_ptr1 + r0, None)
    tmp10 = tl.load(in_ptr1 + (4 + r0), None)
    tmp0 = tl.full([1, 1], 1, tl.int32)
    tmp1 = tmp0 == tmp0
    tmp3 = tl.where(tmp1, tmp2, tmp2)
    tmp4 = tl.broadcast_to(tmp3, [XBLOCK, RBLOCK])
    tmp6 = triton_helpers.max2(tmp4, 1)[:, None]
    tmp7 = tl.full([1, 1], 0, tl.int32)
    tmp8 = tmp0 == tmp7
    tmp11 = tl.where(tmp8, tmp9, tmp10)
    tmp12 = libdevice.isinf(tmp11).to(tl.int1)
    tmp13 = tl.where(tmp12, tmp6, tmp3)
    tl.store(out_ptr2 + tl.broadcast_to(4 + r0, [XBLOCK, RBLOCK]), tmp13, None)


@triton.jit
def triton_poi_fused_index_put_lift_fresh_8(in_ptr0, out_ptr1, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 4
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp3 = tl.load(in_ptr0 + (4 + x0), xmask)
    tmp4 = tl.load(in_ptr0 + (8 + x0), xmask)
    tmp0 = tl.full([1], 2, tl.int32)
    tmp1 = tl.full([1], 1, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp5 = tl.where(tmp2, tmp3, tmp4)
    tmp6 = libdevice.isinf(tmp5).to(tl.int1)
    tmp7 = 0.0
    tmp8 = tl.where(tmp6, tmp7, tmp5)
    tl.store(out_ptr1 + (8 + x0), tmp8, xmask)


@triton.jit
def triton_poi_fused_9(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4
    x0 = xindex % 4
    x2 = xindex
    tmp3 = tl.load(in_ptr0 + (8 + x0), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + x2, xmask)
    tmp0 = x1
    tmp1 = tl.full([1], 2, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp5 = tl.where(tmp2, tmp3, tmp4)
    tl.store(out_ptr0 + x2, tmp5, xmask)


@triton.jit
def triton_per_fused_index_put_max_10(in_ptr0, in_ptr1, out_ptr2, xnumel,
    rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 4
    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
    tmp2 = tl.load(in_ptr0 + (8 + r0), None)
    tmp9 = tl.load(in_ptr1 + (4 + r0), None)
    tmp10 = tl.load(in_ptr1 + (8 + r0), None)
    tmp0 = tl.full([1, 1], 2, tl.int32)
    tmp1 = tmp0 == tmp0
    tmp3 = tl.where(tmp1, tmp2, tmp2)
    tmp4 = tl.broadcast_to(tmp3, [XBLOCK, RBLOCK])
    tmp6 = triton_helpers.max2(tmp4, 1)[:, None]
    tmp7 = tl.full([1, 1], 1, tl.int32)
    tmp8 = tmp0 == tmp7
    tmp11 = tl.where(tmp8, tmp9, tmp10)
    tmp12 = libdevice.isinf(tmp11).to(tl.int1)
    tmp13 = tl.where(tmp12, tmp6, tmp3)
    tl.store(out_ptr2 + tl.broadcast_to(8 + r0, [XBLOCK, RBLOCK]), tmp13, None)


@triton.jit
def triton_poi_fused_index_put_lift_fresh_11(in_ptr0, out_ptr1, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 4
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp3 = tl.load(in_ptr0 + (8 + x0), xmask)
    tmp4 = tl.load(in_ptr0 + (12 + x0), xmask)
    tmp0 = tl.full([1], 3, tl.int32)
    tmp1 = tl.full([1], 2, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp5 = tl.where(tmp2, tmp3, tmp4)
    tmp6 = libdevice.isinf(tmp5).to(tl.int1)
    tmp7 = 0.0
    tmp8 = tl.where(tmp6, tmp7, tmp5)
    tl.store(out_ptr1 + (12 + x0), tmp8, xmask)


@triton.jit
def triton_poi_fused_12(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4
    x0 = xindex % 4
    x2 = xindex
    tmp3 = tl.load(in_ptr0 + (12 + x0), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + x2, xmask)
    tmp0 = x1
    tmp1 = tl.full([1], 3, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp5 = tl.where(tmp2, tmp3, tmp4)
    tl.store(out_ptr0 + x2, tmp5, xmask)


@triton.jit
def triton_per_fused_index_put_max_13(in_ptr0, in_ptr1, out_ptr2, xnumel,
    rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 4
    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
    tmp2 = tl.load(in_ptr0 + (12 + r0), None)
    tmp9 = tl.load(in_ptr1 + (8 + r0), None)
    tmp10 = tl.load(in_ptr1 + (12 + r0), None)
    tmp0 = tl.full([1, 1], 3, tl.int32)
    tmp1 = tmp0 == tmp0
    tmp3 = tl.where(tmp1, tmp2, tmp2)
    tmp4 = tl.broadcast_to(tmp3, [XBLOCK, RBLOCK])
    tmp6 = triton_helpers.max2(tmp4, 1)[:, None]
    tmp7 = tl.full([1, 1], 2, tl.int32)
    tmp8 = tmp0 == tmp7
    tmp11 = tl.where(tmp8, tmp9, tmp10)
    tmp12 = libdevice.isinf(tmp11).to(tl.int1)
    tmp13 = tl.where(tmp12, tmp6, tmp3)
    tl.store(out_ptr2 + tl.broadcast_to(12 + r0, [XBLOCK, RBLOCK]), tmp13, None
        )


@triton.jit
def triton_per_fused_add_div_mean_mul_rsub_sum_14(in_out_ptr0, in_ptr0,
    in_ptr1, in_ptr2, in_ptr3, xnumel, rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 4
    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 + 4 * r0, None, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr1 + 12)
    tmp5 = tl.broadcast_to(tmp4, [XBLOCK, RBLOCK])
    tmp6 = tl.load(in_ptr1 + 4 * r0, None, eviction_policy='evict_last')
    tmp9 = tl.load(in_ptr0 + (1 + 4 * r0), None, eviction_policy='evict_last')
    tmp10 = tl.load(in_ptr1 + 13)
    tmp11 = tl.broadcast_to(tmp10, [XBLOCK, RBLOCK])
    tmp12 = tl.load(in_ptr1 + (1 + 4 * r0), None, eviction_policy='evict_last')
    tmp16 = tl.load(in_ptr0 + (2 + 4 * r0), None, eviction_policy='evict_last')
    tmp17 = tl.load(in_ptr1 + 14)
    tmp18 = tl.broadcast_to(tmp17, [XBLOCK, RBLOCK])
    tmp19 = tl.load(in_ptr1 + (2 + 4 * r0), None, eviction_policy='evict_last')
    tmp23 = tl.load(in_ptr0 + (3 + 4 * r0), None, eviction_policy='evict_last')
    tmp24 = tl.load(in_ptr1 + 15)
    tmp25 = tl.broadcast_to(tmp24, [XBLOCK, RBLOCK])
    tmp26 = tl.load(in_ptr1 + (3 + 4 * r0), None, eviction_policy='evict_last')
    tmp30 = tl.load(in_ptr2 + 4 * r0, None, eviction_policy='evict_last')
    tmp31 = tl.load(in_ptr3 + 4 * r0, None, eviction_policy='evict_last')
    tmp34 = tl.load(in_ptr2 + (1 + 4 * r0), None, eviction_policy='evict_last')
    tmp35 = tl.load(in_ptr3 + (1 + 4 * r0), None, eviction_policy='evict_last')
    tmp39 = tl.load(in_ptr2 + (2 + 4 * r0), None, eviction_policy='evict_last')
    tmp40 = tl.load(in_ptr3 + (2 + 4 * r0), None, eviction_policy='evict_last')
    tmp44 = tl.load(in_ptr2 + (3 + 4 * r0), None, eviction_policy='evict_last')
    tmp45 = tl.load(in_ptr3 + (3 + 4 * r0), None, eviction_policy='evict_last')
    tmp1 = r0
    tmp2 = tl.full([1, 1], 3, tl.int32)
    tmp3 = tmp1 == tmp2
    tmp7 = tl.where(tmp3, tmp5, tmp6)
    tmp8 = tmp0 * tmp7
    tmp13 = tl.where(tmp3, tmp11, tmp12)
    tmp14 = tmp9 * tmp13
    tmp15 = tmp8 + tmp14
    tmp20 = tl.where(tmp3, tmp18, tmp19)
    tmp21 = tmp16 * tmp20
    tmp22 = tmp15 + tmp21
    tmp27 = tl.where(tmp3, tmp25, tmp26)
    tmp28 = tmp23 * tmp27
    tmp29 = tmp22 + tmp28
    tmp32 = tmp30 + tmp31
    tmp33 = tmp32 * tmp7
    tmp36 = tmp34 + tmp35
    tmp37 = tmp36 * tmp13
    tmp38 = tmp33 + tmp37
    tmp41 = tmp39 + tmp40
    tmp42 = tmp41 * tmp20
    tmp43 = tmp38 + tmp42
    tmp46 = tmp44 + tmp45
    tmp47 = tmp46 * tmp27
    tmp48 = tmp43 + tmp47
    tmp49 = 2.0
    tmp50 = tmp29 * tmp49
    tmp51 = 1e-05
    tmp52 = tmp50 + tmp51
    tmp53 = tmp48 + tmp51
    tmp54 = tmp52 / tmp53
    tmp55 = 1.0
    tmp56 = tmp55 - tmp54
    tmp57 = tl.broadcast_to(tmp56, [XBLOCK, RBLOCK])
    tmp59 = tl.sum(tmp57, 1)[:, None]
    tmp60 = 4.0
    tmp61 = tmp59 / tmp60
    tl.debug_barrier()
    tl.store(in_out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp61, 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((4, 4), (4, 1), torch.float32)
        buf1 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        buf29 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        get_raw_stream(0)
        triton_per_fused_mul_sum_0[grid(16)](arg1_1, arg0_1, buf0, buf1,
            buf29, 16, 16, XBLOCK=8, num_warps=2, num_stages=1)
        del arg0_1
        del arg1_1
        buf2 = empty_strided_cuda((4,), (1,), torch.float32)
        triton_poi_fused_index_put_lift_fresh_1[grid(4)](buf1, buf2, 4,
            XBLOCK=4, num_warps=1, num_stages=1)
        buf4 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        triton_poi_fused_mul_reciprocal_2[grid(16)](buf2, buf1, buf4, 16,
            XBLOCK=16, num_warps=1, num_stages=1)
        triton_per_fused_index_put_max_3[grid(1)](buf2, buf1, buf4, 1, 4,
            XBLOCK=1, num_warps=2, num_stages=1)
        del buf2
        buf7 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        triton_poi_fused_4[grid(16)](buf4, buf7, 16, XBLOCK=16, num_warps=1,
            num_stages=1)
        triton_poi_fused_index_put_lift_fresh_5[grid(4)](buf4, buf7, 4,
            XBLOCK=4, num_warps=1, num_stages=1)
        buf11 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        triton_poi_fused_6[grid(16)](buf7, buf11, 16, XBLOCK=16, num_warps=
            1, num_stages=1)
        triton_per_fused_index_put_max_7[grid(1)](buf7, buf4, buf11, 1, 4,
            XBLOCK=1, num_warps=2, num_stages=1)
        buf14 = buf7
        del buf7
        triton_poi_fused_6[grid(16)](buf11, buf14, 16, XBLOCK=16, num_warps
            =1, num_stages=1)
        triton_poi_fused_index_put_lift_fresh_8[grid(4)](buf11, buf14, 4,
            XBLOCK=4, num_warps=1, num_stages=1)
        buf18 = buf4
        del buf4
        triton_poi_fused_9[grid(16)](buf14, buf18, 16, XBLOCK=16, num_warps
            =1, num_stages=1)
        triton_per_fused_index_put_max_10[grid(1)](buf14, buf11, buf18, 1, 
            4, XBLOCK=1, num_warps=2, num_stages=1)
        buf21 = buf14
        del buf14
        triton_poi_fused_9[grid(16)](buf18, buf21, 16, XBLOCK=16, num_warps
            =1, num_stages=1)
        triton_poi_fused_index_put_lift_fresh_11[grid(4)](buf18, buf21, 4,
            XBLOCK=4, num_warps=1, num_stages=1)
        buf25 = buf11
        del buf11
        triton_poi_fused_12[grid(16)](buf21, buf25, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        triton_per_fused_index_put_max_13[grid(1)](buf21, buf18, buf25, 1, 
            4, XBLOCK=1, num_warps=2, num_stages=1)
        del buf18
        del buf21
        buf31 = empty_strided_cuda((), (), torch.float32)
        buf32 = buf31
        del buf31
        triton_per_fused_add_div_mean_mul_rsub_sum_14[grid(1)](buf32, buf0,
            buf25, buf1, buf29, 1, 4, XBLOCK=1, num_warps=2, num_stages=1)
        del buf0
        del buf1
        del buf25
        del buf29
    return buf32,


def issequenceiterable(obj: 'Any') ->bool:
    """
    Determine if the object is an iterable sequence and is not a string.
    """
    if torch.is_tensor(obj):
        return int(obj.dim()) > 0
    return isinstance(obj, collections.abc.Iterable) and not isinstance(obj,
        str)


def ensure_tuple(vals: 'Any') ->Tuple[Any, ...]:
    """
    Returns a tuple of `vals`.
    """
    if not issequenceiterable(vals):
        vals = vals,
    return tuple(vals)


def ensure_tuple_size(tup: 'Any', dim: 'int', pad_val: 'Any'=0) ->Tuple[Any,
    ...]:
    """
    Returns a copy of `tup` with `dim` values by either shortened or padded with `pad_val` as necessary.
    """
    tup = ensure_tuple(tup) + (pad_val,) * dim
    return tuple(tup[:dim])


def one_hot(labels: 'torch.Tensor', num_classes: 'int', dtype:
    'torch.dtype'=torch.float, dim: 'int'=1) ->torch.Tensor:
    """
    For a tensor `labels` of dimensions B1[spatial_dims], return a tensor of dimensions `BN[spatial_dims]`
    for `num_classes` N number of classes.

    Example:

        For every value v = labels[b,1,h,w], the value in the result at [b,v,h,w] will be 1 and all others 0.
        Note that this will include the background label, thus a binary mask should be treated as having 2 classes.
    """
    assert labels.dim() > 0, 'labels should have dim of 1 or more.'
    if labels.ndim < dim + 1:
        shape = ensure_tuple_size(labels.shape, dim + 1, 1)
        labels = labels.reshape(*shape)
    sh = list(labels.shape)
    assert sh[dim
        ] == 1, 'labels should have a channel with length equals to one.'
    sh[dim] = num_classes
    o = torch.zeros(size=sh, dtype=dtype, device=labels.device)
    labels = o.scatter_(dim=dim, index=labels.long(), value=1)
    return labels


class LossReduction(Enum):
    """
    See also:
        - :py:class:`monai.losses.dice.DiceLoss`
        - :py:class:`monai.losses.dice.GeneralizedDiceLoss`
        - :py:class:`monai.losses.focal_loss.FocalLoss`
        - :py:class:`monai.losses.tversky.TverskyLoss`
    """
    NONE = 'none'
    MEAN = 'mean'
    SUM = 'sum'


class Weight(Enum):
    """
    See also: :py:class:`monai.losses.dice.GeneralizedDiceLoss`
    """
    SQUARE = 'square'
    SIMPLE = 'simple'
    UNIFORM = 'uniform'


class GeneralizedDiceLossNew(_Loss):
    """
    Compute the generalised Dice loss defined in:

        Sudre, C. et. al. (2017) Generalised Dice overlap as a deep learning
        loss function for highly unbalanced segmentations. DLMIA 2017.

    Adapted from:
        https://github.com/NifTK/NiftyNet/blob/v0.6.0/niftynet/layer/loss_segmentation.py#L279
    """

    def __init__(self, include_background: 'bool'=True, to_onehot_y: 'bool'
        =False, sigmoid: 'bool'=False, softmax: 'bool'=False, other_act:
        'Optional[Callable]'=None, w_type: 'Union[Weight, str]'=Weight.
        SQUARE, reduction: 'Union[LossReduction, str]'=LossReduction.MEAN,
        smooth_nr: 'float'=1e-05, smooth_dr: 'float'=1e-05, batch: 'bool'=False
        ) ->None:
        """
        Args:
            include_background: If False channel index 0 (background category) is excluded from the calculation.
            to_onehot_y: whether to convert `y` into the one-hot format. Defaults to False.
            sigmoid: If True, apply a sigmoid function to the prediction.
            softmax: If True, apply a softmax function to the prediction.
            other_act: if don't want to use `sigmoid` or `softmax`, use other callable function to execute
                other activation layers, Defaults to ``None``. for example:
                `other_act = torch.tanh`.
            squared_pred: use squared versions of targets and predictions in the denominator or not.
            w_type: {``"square"``, ``"simple"``, ``"uniform"``}
                Type of function to transform ground truth volume to a weight factor. Defaults to ``"square"``.
            reduction: {``"none"``, ``"mean"``, ``"sum"``}
                Specifies the reduction to apply to the output. Defaults to ``"mean"``.

                - ``"none"``: no reduction will be applied.
                - ``"mean"``: the sum of the output will be divided by the number of elements in the output.
                - ``"sum"``: the output will be summed.
            smooth_nr: a small constant added to the numerator to avoid zero.
            smooth_dr: a small constant added to the denominator to avoid nan.
            batch: whether to sum the intersection and union areas over the batch dimension before the dividing.
                Defaults to False, intersection over union is computed from each item in the batch.

        Raises:
            TypeError: When ``other_act`` is not an ``Optional[Callable]``.
            ValueError: When more than 1 of [``sigmoid=True``, ``softmax=True``, ``other_act is not None``].
                Incompatible values.

        """
        super().__init__(reduction=LossReduction(reduction).value)
        if other_act is not None and not callable(other_act):
            raise TypeError(
                f'other_act must be None or callable but is {type(other_act).__name__}.'
                )
        if int(sigmoid) + int(softmax) + int(other_act is not None) > 1:
            raise ValueError(
                'Incompatible values: more than 1 of [sigmoid=True, softmax=True, other_act is not None].'
                )
        self.include_background = include_background
        self.to_onehot_y = to_onehot_y
        self.sigmoid = sigmoid
        self.softmax = softmax
        self.other_act = other_act
        w_type = Weight(w_type)
        self.w_func: 'Callable' = torch.ones_like
        if w_type == Weight.SIMPLE:
            self.w_func = torch.reciprocal
        elif w_type == Weight.SQUARE:
            self.w_func = lambda x: torch.reciprocal(x * x)
        self.smooth_nr = float(smooth_nr)
        self.smooth_dr = float(smooth_dr)
        self.batch = batch

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