# AOT ID: ['0_forward']
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/7e/c7edgnsiuilw7uzwau7radvkvvtmowm7d7uh56mczbhieiykfrnx.py
# Topologically Sorted Source Nodes: [h], Original ATen: [aten.new_zeros]
# Source node to ATen node mapping:
#   h => full_default
# Graph fragment:
#   %full_default : [num_users=3] = call_function[target=torch.ops.aten.full.default](args = ([4, 4], 0), kwargs = {dtype: torch.float32, layout: torch.strided, device: cuda:0, pin_memory: False})
triton_poi_fused_new_zeros_0 = async_compile.triton('triton_poi_fused_new_zeros_0', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

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

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


# kernel path: runs/run_shard_4/inductor_cache/zw/czwxbu6nuxhoujusu3krfetmqzx7rxloioah6gicy4ie2wmv6tqi.py
# Topologically Sorted Source Nodes: [stack_3], Original ATen: [aten.stack]
# Source node to ATen node mapping:
#   stack_3 => cat
# Graph fragment:
#   %cat : [num_users=1] = call_function[target=torch.ops.aten.cat.default](args = ([%primals_8, %primals_8, %primals_8, %primals_8, %primals_6, %primals_6, %primals_6],), kwargs = {})
triton_poi_fused_stack_1 = async_compile.triton('triton_poi_fused_stack_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=[32], 
    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), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_stack_1', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 7, '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_stack_1(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 28
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = (xindex // 4)
    x0 = xindex % 4
    x2 = xindex
    tmp0 = x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 1, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.load(in_ptr0 + (x0), tmp4 & xmask, eviction_policy='evict_last', other=0.0)
    tmp6 = tmp0 >= tmp3
    tmp7 = tl.full([1], 2, tl.int64)
    tmp8 = tmp0 < tmp7
    tmp9 = tmp6 & tmp8
    tmp10 = tl.load(in_ptr0 + (x0), tmp9 & xmask, eviction_policy='evict_last', other=0.0)
    tmp11 = tmp0 >= tmp7
    tmp12 = tl.full([1], 3, tl.int64)
    tmp13 = tmp0 < tmp12
    tmp14 = tmp11 & tmp13
    tmp15 = tl.load(in_ptr0 + (x0), tmp14 & xmask, eviction_policy='evict_last', other=0.0)
    tmp16 = tmp0 >= tmp12
    tmp17 = tl.full([1], 4, tl.int64)
    tmp18 = tmp0 < tmp17
    tmp19 = tmp16 & tmp18
    tmp20 = tl.load(in_ptr0 + (x0), tmp19 & xmask, eviction_policy='evict_last', other=0.0)
    tmp21 = tmp0 >= tmp17
    tmp22 = tl.full([1], 5, tl.int64)
    tmp23 = tmp0 < tmp22
    tmp24 = tmp21 & tmp23
    tmp25 = tl.load(in_ptr1 + (x0), tmp24 & xmask, eviction_policy='evict_last', other=0.0)
    tmp26 = tmp0 >= tmp22
    tmp27 = tl.full([1], 6, tl.int64)
    tmp28 = tmp0 < tmp27
    tmp29 = tmp26 & tmp28
    tmp30 = tl.load(in_ptr1 + (x0), tmp29 & xmask, eviction_policy='evict_last', other=0.0)
    tmp31 = tmp0 >= tmp27
    tmp32 = tl.full([1], 7, tl.int64)
    tmp33 = tmp0 < tmp32
    tmp34 = tl.load(in_ptr1 + (x0), tmp31 & xmask, eviction_policy='evict_last', other=0.0)
    tmp35 = tl.where(tmp29, tmp30, tmp34)
    tmp36 = tl.where(tmp24, tmp25, tmp35)
    tmp37 = tl.where(tmp19, tmp20, tmp36)
    tmp38 = tl.where(tmp14, tmp15, tmp37)
    tmp39 = tl.where(tmp9, tmp10, tmp38)
    tmp40 = tl.where(tmp4, tmp5, tmp39)
    tl.store(out_ptr0 + (x2), tmp40, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_4/inductor_cache/hq/chq2ihz43k5f3canufmvrzz3vy4seaa7ohwwzmv2o6szsqgwtn4r.py
# Topologically Sorted Source Nodes: [stack_4], Original ATen: [aten.stack]
# Source node to ATen node mapping:
#   stack_4 => cat_1
# Graph fragment:
#   %cat_1 : [num_users=1] = call_function[target=torch.ops.aten.cat.default](args = ([%primals_9, %primals_9, %primals_9, %primals_9, %primals_7, %primals_7, %primals_7],), kwargs = {})
triton_poi_fused_stack_2 = async_compile.triton('triton_poi_fused_stack_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=[8], 
    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), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_stack_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_stack_2(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 7
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp5 = tl.load(in_ptr0 + (0))
    tmp6 = tl.broadcast_to(tmp5, [XBLOCK])
    tmp23 = tl.load(in_ptr1 + (0))
    tmp24 = tl.broadcast_to(tmp23, [XBLOCK])
    tmp0 = x0
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 1, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp7 = tmp0 >= tmp3
    tmp8 = tl.full([1], 2, tl.int64)
    tmp9 = tmp0 < tmp8
    tmp10 = tmp7 & tmp9
    tmp11 = tmp0 >= tmp8
    tmp12 = tl.full([1], 3, tl.int64)
    tmp13 = tmp0 < tmp12
    tmp14 = tmp11 & tmp13
    tmp15 = tmp0 >= tmp12
    tmp16 = tl.full([1], 4, tl.int64)
    tmp17 = tmp0 < tmp16
    tmp18 = tmp15 & tmp17
    tmp19 = tmp0 >= tmp16
    tmp20 = tl.full([1], 5, tl.int64)
    tmp21 = tmp0 < tmp20
    tmp22 = tmp19 & tmp21
    tmp25 = tmp0 >= tmp20
    tmp26 = tl.full([1], 6, tl.int64)
    tmp27 = tmp0 < tmp26
    tmp28 = tmp25 & tmp27
    tmp29 = tmp0 >= tmp26
    tmp30 = tl.full([1], 7, tl.int64)
    tmp31 = tmp0 < tmp30
    tmp32 = tl.where(tmp28, tmp24, tmp24)
    tmp33 = tl.where(tmp22, tmp24, tmp32)
    tmp34 = tl.where(tmp18, tmp6, tmp33)
    tmp35 = tl.where(tmp14, tmp6, tmp34)
    tmp36 = tl.where(tmp10, tmp6, tmp35)
    tmp37 = tl.where(tmp4, tmp6, tmp36)
    tl.store(out_ptr0 + (x0), tmp37, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_4/inductor_cache/g2/cg23xreycdv6fxtkdfop75dzuhabmjs3xhrvq22ytdld2qihmoyg.py
# Topologically Sorted Source Nodes: [stack_2], Original ATen: [aten.stack]
# Source node to ATen node mapping:
#   stack_2 => cat_2
# Graph fragment:
#   %cat_2 : [num_users=1] = call_function[target=torch.ops.aten.cat.default](args = ([%getitem, %getitem_2, %getitem_4, %getitem_6, %getitem, %getitem_2, %getitem_4],), kwargs = {})
triton_poi_fused_stack_3 = async_compile.triton('triton_poi_fused_stack_3', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

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

@triton_heuristics.pointwise(
    size_hints=[128], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: '*fp32', 3: '*fp32', 4: '*fp32', 5: '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), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_stack_3', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 7, '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_stack_3(in_ptr0, in_ptr1, in_ptr2, in_ptr3, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 112
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = (xindex // 4)
    x0 = xindex % 4
    x2 = xindex
    tmp0 = x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.load(in_ptr0 + (x0 + (4*x1)), tmp4 & xmask, other=0.0)
    tmp6 = tmp0 >= tmp3
    tmp7 = tl.full([1], 8, tl.int64)
    tmp8 = tmp0 < tmp7
    tmp9 = tmp6 & tmp8
    tmp10 = tl.load(in_ptr1 + (x0 + (4*((-4) + x1))), tmp9 & xmask, other=0.0)
    tmp11 = tmp0 >= tmp7
    tmp12 = tl.full([1], 12, tl.int64)
    tmp13 = tmp0 < tmp12
    tmp14 = tmp11 & tmp13
    tmp15 = tl.load(in_ptr2 + (x0 + (4*((-8) + x1))), tmp14 & xmask, other=0.0)
    tmp16 = tmp0 >= tmp12
    tmp17 = tl.full([1], 16, tl.int64)
    tmp18 = tmp0 < tmp17
    tmp19 = tmp16 & tmp18
    tmp20 = tl.load(in_ptr3 + (x0 + (4*((-12) + x1))), tmp19 & xmask, other=0.0)
    tmp21 = tmp0 >= tmp17
    tmp22 = tl.full([1], 20, tl.int64)
    tmp23 = tmp0 < tmp22
    tmp24 = tmp21 & tmp23
    tmp25 = tl.load(in_ptr0 + (x0 + (4*((-16) + x1))), tmp24 & xmask, other=0.0)
    tmp26 = tmp0 >= tmp22
    tmp27 = tl.full([1], 24, tl.int64)
    tmp28 = tmp0 < tmp27
    tmp29 = tmp26 & tmp28
    tmp30 = tl.load(in_ptr1 + (x0 + (4*((-20) + x1))), tmp29 & xmask, other=0.0)
    tmp31 = tmp0 >= tmp27
    tmp32 = tl.full([1], 28, tl.int64)
    tmp33 = tmp0 < tmp32
    tmp34 = tl.load(in_ptr2 + (x0 + (4*((-24) + x1))), tmp31 & xmask, other=0.0)
    tmp35 = tl.where(tmp29, tmp30, tmp34)
    tmp36 = tl.where(tmp24, tmp25, tmp35)
    tmp37 = tl.where(tmp19, tmp20, tmp36)
    tmp38 = tl.where(tmp14, tmp15, tmp37)
    tmp39 = tl.where(tmp9, tmp10, tmp38)
    tmp40 = tl.where(tmp4, tmp5, tmp39)
    tl.store(out_ptr0 + (x2), tmp40, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_4/inductor_cache/oo/coo3437vo4xsrwedl7lcrwfswym3k6uq6mzbtutpnydbupg5rgob.py
# Topologically Sorted Source Nodes: [p_m, sigmoid_2, sub_8, bernoulli_2, sigmoid_1, sub_4, bernoulli_1, sigmoid, un_halted_prob_1, p_n_1, un_halted_prob_2, p_n_2, p_n_3, bernoulli, halted_1, sub_2, mul_3, mul_4, p_m_1, halt_1, sub_6, mul_10, mul_11, p_m_2, halted_2, sub_9, halt_2, sub_10, mul_17, mul_18, p_m_3, halted_3, mul_20, mul_13, mul_6, y_m_1, mul_12, y_m_2, mul_19, y_m_3, lambda_n_3, bernoulli_3, sub_13, halt_3, sub_14, mul_24, mul_25, p_m_4, mul_26, mul_27, y_m_4], Original ATen: [aten.new_zeros, aten.sigmoid, aten.rsub, aten.bernoulli, aten.mul, aten.add, aten.new_ones]
# Source node to ATen node mapping:
#   bernoulli => lt_2
#   bernoulli_1 => convert_element_type_1, lt_1
#   bernoulli_2 => convert_element_type, lt
#   bernoulli_3 => convert_element_type_3, lt_3
#   halt_1 => mul_10
#   halt_2 => mul_13
#   halt_3 => mul_23
#   halted_1 => convert_element_type_2
#   halted_2 => add_4
#   halted_3 => add_6
#   lambda_n_3 => full_default_6
#   mul_10 => mul_11
#   mul_11 => mul_12
#   mul_12 => mul_19
#   mul_13 => mul_17
#   mul_17 => mul_14
#   mul_18 => mul_15
#   mul_19 => mul_20
#   mul_20 => mul_16
#   mul_24 => mul_24
#   mul_25 => mul_25
#   mul_26 => mul_26
#   mul_27 => mul_27
#   mul_3 => mul_7
#   mul_4 => mul_8
#   mul_6 => mul_18
#   p_m => full_default_3
#   p_m_1 => add_1
#   p_m_2 => add_3
#   p_m_3 => add_5
#   p_m_4 => add_10
#   p_n_1 => mul_2
#   p_n_2 => mul_4
#   p_n_3 => mul_5
#   sigmoid => sigmoid_2
#   sigmoid_1 => sigmoid_1
#   sigmoid_2 => sigmoid
#   sub_10 => sub_10
#   sub_13 => sub_13
#   sub_14 => sub_14
#   sub_2 => sub_4
#   sub_4 => sub_2
#   sub_6 => sub_7
#   sub_8 => sub_1
#   sub_9 => sub_9
#   un_halted_prob_1 => sub_3
#   un_halted_prob_2 => mul_3
#   y_m_1 => add_7
#   y_m_2 => add_8
#   y_m_3 => add_9
#   y_m_4 => add_11
# Graph fragment:
#   %full_default_3 : [num_users=1] = call_function[target=torch.ops.aten.full.default](args = ([4], 0), kwargs = {dtype: torch.float32, layout: torch.strided, device: cuda:0, pin_memory: False})
#   %sigmoid : [num_users=2] = call_function[target=torch.ops.aten.sigmoid.default](args = (%getitem_14,), kwargs = {})
#   %sub_1 : [num_users=1] = call_function[target=torch.ops.aten.sub.Tensor](args = (1, %select), kwargs = {})
#   %lt : [num_users=2] = call_function[target=torch.ops.aten.lt.Tensor](args = (%rand, %select), kwargs = {})
#   %convert_element_type : [num_users=1] = call_function[target=torch.ops.prims.convert_element_type.default](args = (%lt, torch.float32), kwargs = {})
#   %sigmoid_1 : [num_users=2] = call_function[target=torch.ops.aten.sigmoid.default](args = (%getitem_13,), kwargs = {})
#   %sub_2 : [num_users=1] = call_function[target=torch.ops.aten.sub.Tensor](args = (1, %select_1), kwargs = {})
#   %lt_1 : [num_users=2] = call_function[target=torch.ops.aten.lt.Tensor](args = (%rand_1, %select_1), kwargs = {})
#   %convert_element_type_1 : [num_users=1] = call_function[target=torch.ops.prims.convert_element_type.default](args = (%lt_1, torch.float32), kwargs = {})
#   %sigmoid_2 : [num_users=2] = call_function[target=torch.ops.aten.sigmoid.default](args = (%getitem_12,), kwargs = {})
#   %sub_3 : [num_users=2] = call_function[target=torch.ops.aten.sub.Tensor](args = (1, %select_2), kwargs = {})
#   %mul_2 : [num_users=2] = call_function[target=torch.ops.aten.mul.Tensor](args = (%sub_3, %select_1), kwargs = {})
#   %mul_3 : [num_users=2] = call_function[target=torch.ops.aten.mul.Tensor](args = (%sub_3, %sub_2), kwargs = {})
#   %mul_4 : [num_users=2] = call_function[target=torch.ops.aten.mul.Tensor](args = (%mul_3, %select), kwargs = {})
#   %mul_5 : [num_users=2] = call_function[target=torch.ops.aten.mul.Tensor](args = (%mul_3, %sub_1), kwargs = {})
#   %lt_2 : [num_users=2] = call_function[target=torch.ops.aten.lt.Tensor](args = (%rand_2, %select_2), kwargs = {})
#   %convert_element_type_2 : [num_users=4] = call_function[target=torch.ops.prims.convert_element_type.default](args = (%lt_2, torch.float32), kwargs = {})
#   %sub_4 : [num_users=2] = call_function[target=torch.ops.aten.sub.Tensor](args = (1, %convert_element_type_2), kwargs = {})
#   %mul_7 : [num_users=2] = call_function[target=torch.ops.aten.mul.Tensor](args = (%full_default_3, %sub_4), kwargs = {})
#   %mul_8 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%select_2, %convert_element_type_2), kwargs = {})
#   %add_1 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul_7, %mul_8), kwargs = {})
#   %mul_10 : [num_users=4] = call_function[target=torch.ops.aten.mul.Tensor](args = (%convert_element_type_1, %sub_4), kwargs = {})
#   %sub_7 : [num_users=2] = call_function[target=torch.ops.aten.sub.Tensor](args = (1, %mul_10), kwargs = {})
#   %mul_11 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%add_1, %sub_7), kwargs = {})
#   %mul_12 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%mul_2, %mul_10), kwargs = {})
#   %add_3 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul_11, %mul_12), kwargs = {})
#   %add_4 : [num_users=2] = call_function[target=torch.ops.aten.add.Tensor](args = (%convert_element_type_2, %mul_10), kwargs = {})
#   %sub_9 : [num_users=1] = call_function[target=torch.ops.aten.sub.Tensor](args = (1, %add_4), kwargs = {})
#   %mul_13 : [num_users=4] = call_function[target=torch.ops.aten.mul.Tensor](args = (%convert_element_type, %sub_9), kwargs = {})
#   %sub_10 : [num_users=2] = call_function[target=torch.ops.aten.sub.Tensor](args = (1, %mul_13), kwargs = {})
#   %mul_14 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%add_3, %sub_10), kwargs = {})
#   %mul_15 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%mul_4, %mul_13), kwargs = {})
#   %add_5 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul_14, %mul_15), kwargs = {})
#   %add_6 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%add_4, %mul_13), kwargs = {})
#   %mul_16 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%select_3, %mul_13), kwargs = {})
#   %mul_17 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%select_4, %mul_10), kwargs = {})
#   %mul_18 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%select_5, %convert_element_type_2), kwargs = {})
#   %add_7 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul_7, %mul_18), kwargs = {})
#   %mul_19 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%add_7, %sub_7), kwargs = {})
#   %add_8 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul_19, %mul_17), kwargs = {})
#   %mul_20 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%add_8, %sub_10), kwargs = {})
#   %add_9 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul_20, %mul_16), kwargs = {})
#   %full_default_6 : [num_users=1] = call_function[target=torch.ops.aten.full.default](args = ([4], 1), kwargs = {dtype: torch.float32, layout: torch.strided, device: cuda:0, pin_memory: False})
#   %lt_3 : [num_users=2] = call_function[target=torch.ops.aten.lt.Tensor](args = (%rand_3, %full_default_6), kwargs = {})
#   %convert_element_type_3 : [num_users=1] = call_function[target=torch.ops.prims.convert_element_type.default](args = (%lt_3, torch.float32), kwargs = {})
#   %sub_13 : [num_users=1] = call_function[target=torch.ops.aten.sub.Tensor](args = (1, %add_6), kwargs = {})
#   %mul_23 : [num_users=3] = call_function[target=torch.ops.aten.mul.Tensor](args = (%convert_element_type_3, %sub_13), kwargs = {})
#   %sub_14 : [num_users=2] = call_function[target=torch.ops.aten.sub.Tensor](args = (1, %mul_23), kwargs = {})
#   %mul_24 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%add_5, %sub_14), kwargs = {})
#   %mul_25 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%mul_5, %mul_23), kwargs = {})
#   %add_10 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul_24, %mul_25), kwargs = {})
#   %mul_26 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%add_9, %sub_14), kwargs = {})
#   %mul_27 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%select_6, %mul_23), kwargs = {})
#   %add_11 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul_26, %mul_27), kwargs = {})
triton_poi_fused_add_bernoulli_mul_new_ones_new_zeros_rsub_sigmoid_4 = async_compile.triton('triton_poi_fused_add_bernoulli_mul_new_ones_new_zeros_rsub_sigmoid_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=[4], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: '*fp32', 3: '*fp32', 4: '*fp32', 5: '*fp32', 6: '*fp32', 7: '*fp32', 8: '*i1', 9: '*fp32', 10: '*fp32', 11: '*fp32', 12: '*i1', 13: '*i1', 14: '*i1', 15: '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, 7, 8, 9, 10, 11, 12, 13, 14), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_bernoulli_mul_new_ones_new_zeros_rsub_sigmoid_4', 'mutated_arg_names': ['in_out_ptr0', 'in_out_ptr1'], 'no_x_dim': False, 'num_load': 18, '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_add_bernoulli_mul_new_ones_new_zeros_rsub_sigmoid_4(in_out_ptr0, in_out_ptr1, in_ptr0, in_ptr1, in_ptr2, in_ptr3, in_ptr4, in_ptr5, out_ptr0, out_ptr1, out_ptr2, out_ptr3, out_ptr4, out_ptr5, out_ptr6, 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)
    tmp3 = tl.load(in_ptr1 + (4))
    tmp4 = tl.broadcast_to(tmp3, [XBLOCK])
    tmp5 = tl.load(in_ptr2 + (16 + x0), xmask)
    tmp8 = tl.load(in_ptr1 + (5))
    tmp9 = tl.broadcast_to(tmp8, [XBLOCK])
    tmp10 = tl.load(in_ptr2 + (20 + x0), xmask)
    tmp13 = tl.load(in_ptr1 + (6))
    tmp14 = tl.broadcast_to(tmp13, [XBLOCK])
    tmp15 = tl.load(in_ptr2 + (24 + x0), xmask)
    tmp18 = tl.load(in_ptr3 + (x0), xmask)
    tmp20 = tl.load(in_ptr4 + (x0), xmask)
    tmp22 = tl.load(in_ptr5 + (x0), xmask)
    tmp59 = tl.load(in_ptr1 + (0))
    tmp60 = tl.broadcast_to(tmp59, [XBLOCK])
    tmp61 = tl.load(in_ptr2 + (x0), xmask)
    tmp66 = tl.load(in_ptr1 + (1))
    tmp67 = tl.broadcast_to(tmp66, [XBLOCK])
    tmp68 = tl.load(in_ptr2 + (4 + x0), xmask)
    tmp73 = tl.load(in_ptr1 + (2))
    tmp74 = tl.broadcast_to(tmp73, [XBLOCK])
    tmp75 = tl.load(in_ptr2 + (8 + x0), xmask)
    tmp80 = tl.load(in_ptr1 + (3))
    tmp81 = tl.broadcast_to(tmp80, [XBLOCK])
    tmp82 = tl.load(in_ptr2 + (12 + x0), xmask)
    tmp1 = 1.0
    tmp2 = tmp0 < tmp1
    tmp6 = tmp4 + tmp5
    tmp7 = tl.sigmoid(tmp6)
    tmp11 = tmp9 + tmp10
    tmp12 = tl.sigmoid(tmp11)
    tmp16 = tmp14 + tmp15
    tmp17 = tl.sigmoid(tmp16)
    tmp19 = tmp18 < tmp17
    tmp21 = tmp20 < tmp12
    tmp23 = tmp22 < tmp7
    tmp24 = tmp23.to(tl.float32)
    tmp25 = tmp1 - tmp24
    tmp26 = 0.0
    tmp27 = tmp26 * tmp25
    tmp28 = tmp7 * tmp24
    tmp29 = tmp27 + tmp28
    tmp30 = tmp21.to(tl.float32)
    tmp31 = tmp30 * tmp25
    tmp32 = tmp1 - tmp31
    tmp33 = tmp29 * tmp32
    tmp34 = tmp1 - tmp7
    tmp35 = tmp34 * tmp12
    tmp36 = tmp35 * tmp31
    tmp37 = tmp33 + tmp36
    tmp38 = tmp19.to(tl.float32)
    tmp39 = tmp24 + tmp31
    tmp40 = tmp1 - tmp39
    tmp41 = tmp38 * tmp40
    tmp42 = tmp1 - tmp41
    tmp43 = tmp37 * tmp42
    tmp44 = tmp1 - tmp12
    tmp45 = tmp34 * tmp44
    tmp46 = tmp45 * tmp17
    tmp47 = tmp46 * tmp41
    tmp48 = tmp43 + tmp47
    tmp49 = tmp2.to(tl.float32)
    tmp50 = tmp39 + tmp41
    tmp51 = tmp1 - tmp50
    tmp52 = tmp49 * tmp51
    tmp53 = tmp1 - tmp52
    tmp54 = tmp48 * tmp53
    tmp55 = tmp1 - tmp17
    tmp56 = tmp45 * tmp55
    tmp57 = tmp56 * tmp52
    tmp58 = tmp54 + tmp57
    tmp62 = tmp60 + tmp61
    tmp63 = tmp62 * tmp24
    tmp64 = tmp27 + tmp63
    tmp65 = tmp64 * tmp32
    tmp69 = tmp67 + tmp68
    tmp70 = tmp69 * tmp31
    tmp71 = tmp65 + tmp70
    tmp72 = tmp71 * tmp42
    tmp76 = tmp74 + tmp75
    tmp77 = tmp76 * tmp41
    tmp78 = tmp72 + tmp77
    tmp79 = tmp78 * tmp53
    tmp83 = tmp81 + tmp82
    tmp84 = tmp83 * tmp52
    tmp85 = tmp79 + tmp84
    tl.store(out_ptr0 + (x0), tmp2, xmask)
    tl.store(out_ptr1 + (x0), tmp7, xmask)
    tl.store(out_ptr2 + (x0), tmp12, xmask)
    tl.store(out_ptr3 + (x0), tmp17, xmask)
    tl.store(out_ptr4 + (x0), tmp19, xmask)
    tl.store(out_ptr5 + (x0), tmp21, xmask)
    tl.store(out_ptr6 + (x0), tmp23, xmask)
    tl.store(in_out_ptr0 + (x0), tmp58, xmask)
    tl.store(in_out_ptr1 + (x0), tmp85, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_4/inductor_cache/jy/cjys3iopkrrusczes7dimlrpecwecndhklisdis5ynngdb4dpi3f.py
# Topologically Sorted Source Nodes: [stack], Original ATen: [aten.stack]
# Source node to ATen node mapping:
#   stack => cat_3
# Graph fragment:
#   %cat_3 : [num_users=1] = call_function[target=torch.ops.aten.cat.default](args = ([%select_2, %mul_2, %mul_4, %mul_5],), kwargs = {})
triton_poi_fused_stack_5 = async_compile.triton('triton_poi_fused_stack_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=[16], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: '*fp32', 3: '*fp32', 4: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2, 3, 4), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_stack_5', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 9, '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_stack_5(in_ptr0, in_ptr1, in_ptr2, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = x0
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.load(in_ptr0 + (x0), tmp4 & xmask, eviction_policy='evict_last', other=0.0)
    tmp6 = tmp0 >= tmp3
    tmp7 = tl.full([1], 8, tl.int64)
    tmp8 = tmp0 < tmp7
    tmp9 = tmp6 & tmp8
    tmp10 = tl.load(in_ptr0 + ((-4) + x0), tmp9 & xmask, eviction_policy='evict_last', other=0.0)
    tmp11 = 1.0
    tmp12 = tmp11 - tmp10
    tmp13 = tl.load(in_ptr1 + ((-4) + x0), tmp9 & xmask, eviction_policy='evict_last', other=0.0)
    tmp14 = tmp12 * tmp13
    tmp15 = tl.full(tmp14.shape, 0.0, tmp14.dtype)
    tmp16 = tl.where(tmp9, tmp14, tmp15)
    tmp17 = tmp0 >= tmp7
    tmp18 = tl.full([1], 12, tl.int64)
    tmp19 = tmp0 < tmp18
    tmp20 = tmp17 & tmp19
    tmp21 = tl.load(in_ptr0 + ((-8) + x0), tmp20 & xmask, eviction_policy='evict_last', other=0.0)
    tmp22 = tmp11 - tmp21
    tmp23 = tl.load(in_ptr1 + ((-8) + x0), tmp20 & xmask, eviction_policy='evict_last', other=0.0)
    tmp24 = tmp11 - tmp23
    tmp25 = tmp22 * tmp24
    tmp26 = tl.load(in_ptr2 + ((-8) + x0), tmp20 & xmask, eviction_policy='evict_last', other=0.0)
    tmp27 = tmp25 * tmp26
    tmp28 = tl.full(tmp27.shape, 0.0, tmp27.dtype)
    tmp29 = tl.where(tmp20, tmp27, tmp28)
    tmp30 = tmp0 >= tmp18
    tmp31 = tl.full([1], 16, tl.int64)
    tmp32 = tmp0 < tmp31
    tmp33 = tl.load(in_ptr0 + ((-12) + x0), tmp30 & xmask, eviction_policy='evict_last', other=0.0)
    tmp34 = tmp11 - tmp33
    tmp35 = tl.load(in_ptr1 + ((-12) + x0), tmp30 & xmask, eviction_policy='evict_last', other=0.0)
    tmp36 = tmp11 - tmp35
    tmp37 = tmp34 * tmp36
    tmp38 = tl.load(in_ptr2 + ((-12) + x0), tmp30 & xmask, eviction_policy='evict_last', other=0.0)
    tmp39 = tmp11 - tmp38
    tmp40 = tmp37 * tmp39
    tmp41 = tl.full(tmp40.shape, 0.0, tmp40.dtype)
    tmp42 = tl.where(tmp30, tmp40, tmp41)
    tmp43 = tl.where(tmp20, tmp29, tmp42)
    tmp44 = tl.where(tmp9, tmp16, tmp43)
    tmp45 = tl.where(tmp4, tmp5, tmp44)
    tl.store(out_ptr0 + (x0), tmp45, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_4/inductor_cache/o4/co4faoo6vn5p4s7le2cjisym3ezxxx6wg265nankhpf5szhdz3rv.py
# Topologically Sorted Source Nodes: [stack_1], Original ATen: [aten.stack]
# Source node to ATen node mapping:
#   stack_1 => cat_4
# Graph fragment:
#   %cat_4 : [num_users=1] = call_function[target=torch.ops.aten.cat.default](args = ([%select_5, %select_4, %select_3, %select_6],), kwargs = {})
triton_poi_fused_stack_6 = async_compile.triton('triton_poi_fused_stack_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: '*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_stack_6', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 8, 'num_reduction': 0, 'backend_hash': 'A9C866B4A14FD3277824029365D703C2427B2E685E54EC9B3EF4ADC8D1EEAC1D', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
    min_elem_per_thread=0
)
@triton.jit
def triton_poi_fused_stack_6(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
    x0 = xindex
    tmp5 = tl.load(in_ptr0 + (0))
    tmp6 = tl.broadcast_to(tmp5, [XBLOCK])
    tmp15 = tl.load(in_ptr0 + (1))
    tmp16 = tl.broadcast_to(tmp15, [XBLOCK])
    tmp25 = tl.load(in_ptr0 + (2))
    tmp26 = tl.broadcast_to(tmp25, [XBLOCK])
    tmp34 = tl.load(in_ptr0 + (3))
    tmp35 = tl.broadcast_to(tmp34, [XBLOCK])
    tmp0 = x0
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp7 = tl.load(in_ptr1 + (x0), tmp4 & xmask, eviction_policy='evict_last', other=0.0)
    tmp8 = tmp6 + tmp7
    tmp9 = tl.full(tmp8.shape, 0.0, tmp8.dtype)
    tmp10 = tl.where(tmp4, tmp8, tmp9)
    tmp11 = tmp0 >= tmp3
    tmp12 = tl.full([1], 8, tl.int64)
    tmp13 = tmp0 < tmp12
    tmp14 = tmp11 & tmp13
    tmp17 = tl.load(in_ptr1 + (4 + ((-4) + x0)), tmp14 & xmask, eviction_policy='evict_last', other=0.0)
    tmp18 = tmp16 + tmp17
    tmp19 = tl.full(tmp18.shape, 0.0, tmp18.dtype)
    tmp20 = tl.where(tmp14, tmp18, tmp19)
    tmp21 = tmp0 >= tmp12
    tmp22 = tl.full([1], 12, tl.int64)
    tmp23 = tmp0 < tmp22
    tmp24 = tmp21 & tmp23
    tmp27 = tl.load(in_ptr1 + (8 + ((-8) + x0)), tmp24 & xmask, eviction_policy='evict_last', other=0.0)
    tmp28 = tmp26 + tmp27
    tmp29 = tl.full(tmp28.shape, 0.0, tmp28.dtype)
    tmp30 = tl.where(tmp24, tmp28, tmp29)
    tmp31 = tmp0 >= tmp22
    tmp32 = tl.full([1], 16, tl.int64)
    tmp33 = tmp0 < tmp32
    tmp36 = tl.load(in_ptr1 + (12 + ((-12) + x0)), tmp31 & xmask, eviction_policy='evict_last', other=0.0)
    tmp37 = tmp35 + tmp36
    tmp38 = tl.full(tmp37.shape, 0.0, tmp37.dtype)
    tmp39 = tl.where(tmp31, tmp37, tmp38)
    tmp40 = tl.where(tmp24, tmp30, tmp39)
    tmp41 = tl.where(tmp14, tmp20, tmp40)
    tmp42 = tl.where(tmp4, tmp10, tmp41)
    tl.store(out_ptr0 + (x0), tmp42, xmask)
''', device_str='cuda')


async_compile.wait(globals())
del async_compile

def call(args):
    primals_1, primals_2, primals_3, primals_4, primals_5, primals_6, primals_7, primals_8, primals_9 = args
    args.clear()
    assert_size_stride(primals_1, (4, 4), (4, 1))
    assert_size_stride(primals_2, (12, 4), (4, 1))
    assert_size_stride(primals_3, (12, 4), (4, 1))
    assert_size_stride(primals_4, (12, ), (1, ))
    assert_size_stride(primals_5, (12, ), (1, ))
    assert_size_stride(primals_6, (1, 4), (4, 1))
    assert_size_stride(primals_7, (1, ), (1, ))
    assert_size_stride(primals_8, (1, 4), (4, 1))
    assert_size_stride(primals_9, (1, ), (1, ))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        # Topologically Sorted Source Nodes: [h], Original ATen: [aten.new_zeros]
        stream0 = get_raw_stream(0)
        triton_poi_fused_new_zeros_0.run(buf0, 16, grid=grid(16), stream=stream0)
        buf1 = empty_strided_cuda((4, 12), (12, 1), torch.float32)
        # Topologically Sorted Source Nodes: [ret], Original ATen: [aten.mm]
        extern_kernels.mm(primals_1, reinterpret_tensor(primals_2, (4, 12), (1, 4), 0), out=buf1)
        del primals_2
        buf2 = empty_strided_cuda((4, 12), (12, 1), torch.float32)
        # Topologically Sorted Source Nodes: [ret], Original ATen: [aten.mm]
        extern_kernels.mm(buf0, reinterpret_tensor(primals_3, (4, 12), (1, 4), 0), out=buf2)
        # Topologically Sorted Source Nodes: [ret], Original ATen: [aten._thnn_fused_gru_cell]
        buf3 = torch.ops.aten._thnn_fused_gru_cell.default(buf1, buf2, buf0, primals_4, primals_5)
        buf4 = buf3[0]
        buf5 = buf3[1]
        del buf3
        buf6 = empty_strided_cuda((7, 4), (4, 1), torch.float32)
        # Topologically Sorted Source Nodes: [stack_3], Original ATen: [aten.stack]
        triton_poi_fused_stack_1.run(primals_8, primals_6, buf6, 28, grid=grid(28), stream=stream0)
        del primals_6
        del primals_8
        buf7 = empty_strided_cuda((7, ), (1, ), torch.float32)
        # Topologically Sorted Source Nodes: [stack_4], Original ATen: [aten.stack]
        triton_poi_fused_stack_2.run(primals_9, primals_7, buf7, 7, grid=grid(7), stream=stream0)
        del primals_7
        del primals_9
        buf8 = buf2; del buf2  # reuse
        # Topologically Sorted Source Nodes: [ret_1], Original ATen: [aten.mm]
        extern_kernels.mm(buf4, reinterpret_tensor(primals_3, (4, 12), (1, 4), 0), out=buf8)
        # Topologically Sorted Source Nodes: [ret_1], Original ATen: [aten._thnn_fused_gru_cell]
        buf9 = torch.ops.aten._thnn_fused_gru_cell.default(buf1, buf8, buf4, primals_4, primals_5)
        buf10 = buf9[0]
        buf11 = buf9[1]
        del buf9
        buf12 = buf8; del buf8  # reuse
        # Topologically Sorted Source Nodes: [ret_2], Original ATen: [aten.mm]
        extern_kernels.mm(buf10, reinterpret_tensor(primals_3, (4, 12), (1, 4), 0), out=buf12)
        # Topologically Sorted Source Nodes: [ret_2], Original ATen: [aten._thnn_fused_gru_cell]
        buf13 = torch.ops.aten._thnn_fused_gru_cell.default(buf1, buf12, buf10, primals_4, primals_5)
        buf14 = buf13[0]
        buf15 = buf13[1]
        del buf13
        buf16 = buf12; del buf12  # reuse
        # Topologically Sorted Source Nodes: [ret_3], Original ATen: [aten.mm]
        extern_kernels.mm(buf14, reinterpret_tensor(primals_3, (4, 12), (1, 4), 0), out=buf16)
        # Topologically Sorted Source Nodes: [ret_3], Original ATen: [aten._thnn_fused_gru_cell]
        buf17 = torch.ops.aten._thnn_fused_gru_cell.default(buf1, buf16, buf14, primals_4, primals_5)
        del buf1
        del buf16
        del primals_4
        del primals_5
        buf18 = buf17[0]
        buf19 = buf17[1]
        del buf17
        buf20 = empty_strided_cuda((28, 4), (4, 1), torch.float32)
        # Topologically Sorted Source Nodes: [stack_2], Original ATen: [aten.stack]
        triton_poi_fused_stack_3.run(buf4, buf10, buf14, buf18, buf20, 112, grid=grid(112), stream=stream0)
        buf21 = empty_strided_cuda((7, 4, 1), (4, 1, 1), torch.float32)
        # Topologically Sorted Source Nodes: [baddbmm], Original ATen: [aten.baddbmm]
        extern_kernels.bmm(reinterpret_tensor(buf20, (7, 4, 4), (16, 4, 1), 0), reinterpret_tensor(buf6, (7, 4, 1), (4, 1, 4), 0), out=buf21)
        # Topologically Sorted Source Nodes: [bernoulli_2], Original ATen: [aten.bernoulli]
        buf23 = torch.ops.aten.rand.default([4], dtype=torch.float32, device=device(type='cuda', index=0), pin_memory=False)
        buf24 = buf23
        del buf23
        # Topologically Sorted Source Nodes: [bernoulli_1], Original ATen: [aten.bernoulli]
        buf27 = torch.ops.aten.rand.default([4], dtype=torch.float32, device=device(type='cuda', index=0), pin_memory=False)
        buf28 = buf27
        del buf27
        # Topologically Sorted Source Nodes: [bernoulli], Original ATen: [aten.bernoulli]
        buf31 = torch.ops.aten.rand.default([4], dtype=torch.float32, device=device(type='cuda', index=0), pin_memory=False)
        buf32 = buf31
        del buf31
        # Topologically Sorted Source Nodes: [bernoulli_3], Original ATen: [aten.bernoulli]
        buf36 = torch.ops.aten.rand.default([4], dtype=torch.float32, device=device(type='cuda', index=0), pin_memory=False)
        buf37 = buf36
        del buf36
        buf38 = empty_strided_cuda((4, ), (1, ), torch.bool)
        buf30 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        buf26 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        buf22 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        buf25 = empty_strided_cuda((4, ), (1, ), torch.bool)
        buf29 = empty_strided_cuda((4, ), (1, ), torch.bool)
        buf33 = empty_strided_cuda((4, ), (1, ), torch.bool)
        buf34 = empty_strided_cuda((4, ), (1, ), torch.float32)
        buf39 = buf34; del buf34  # reuse
        buf35 = empty_strided_cuda((4, ), (1, ), torch.float32)
        buf40 = buf35; del buf35  # reuse
        # Topologically Sorted Source Nodes: [p_m, sigmoid_2, sub_8, bernoulli_2, sigmoid_1, sub_4, bernoulli_1, sigmoid, un_halted_prob_1, p_n_1, un_halted_prob_2, p_n_2, p_n_3, bernoulli, halted_1, sub_2, mul_3, mul_4, p_m_1, halt_1, sub_6, mul_10, mul_11, p_m_2, halted_2, sub_9, halt_2, sub_10, mul_17, mul_18, p_m_3, halted_3, mul_20, mul_13, mul_6, y_m_1, mul_12, y_m_2, mul_19, y_m_3, lambda_n_3, bernoulli_3, sub_13, halt_3, sub_14, mul_24, mul_25, p_m_4, mul_26, mul_27, y_m_4], Original ATen: [aten.new_zeros, aten.sigmoid, aten.rsub, aten.bernoulli, aten.mul, aten.add, aten.new_ones]
        triton_poi_fused_add_bernoulli_mul_new_ones_new_zeros_rsub_sigmoid_4.run(buf39, buf40, buf37, buf7, buf21, buf24, buf28, buf32, buf38, buf30, buf26, buf22, buf25, buf29, buf33, 4, grid=grid(4), stream=stream0)
        del buf24
        del buf28
        del buf32
        del buf37
        buf41 = reinterpret_tensor(buf18, (16, ), (1, ), 0); del buf18  # reuse
        # Topologically Sorted Source Nodes: [stack], Original ATen: [aten.stack]
        triton_poi_fused_stack_5.run(buf30, buf26, buf22, buf41, 16, grid=grid(16), stream=stream0)
        buf42 = empty_strided_cuda((16, ), (1, ), torch.float32)
        # Topologically Sorted Source Nodes: [stack_1], Original ATen: [aten.stack]
        triton_poi_fused_stack_6.run(buf7, buf21, buf42, 16, grid=grid(16), stream=stream0)
        del buf21
        del buf7
    return (reinterpret_tensor(buf41, (4, 4), (4, 1), 0), reinterpret_tensor(buf42, (4, 4), (4, 1), 0), buf39, buf40, primals_1, buf0, buf4, buf5, buf10, buf11, buf14, buf15, buf19, buf22, buf25, buf26, buf29, buf30, buf33, buf38, reinterpret_tensor(buf6, (7, 1, 4), (4, 4, 1), 0), reinterpret_tensor(buf20, (7, 4, 4), (16, 1, 4), 0), primals_3, )


def benchmark_compiled_module(times=10, repeat=10):
    from torch._dynamo.testing import rand_strided
    from torch._inductor.utils import print_performance
    primals_1 = rand_strided((4, 4), (4, 1), device='cuda:0', dtype=torch.float32)
    primals_2 = rand_strided((12, 4), (4, 1), device='cuda:0', dtype=torch.float32)
    primals_3 = rand_strided((12, 4), (4, 1), device='cuda:0', dtype=torch.float32)
    primals_4 = rand_strided((12, ), (1, ), device='cuda:0', dtype=torch.float32)
    primals_5 = rand_strided((12, ), (1, ), device='cuda:0', dtype=torch.float32)
    primals_6 = rand_strided((1, 4), (4, 1), device='cuda:0', dtype=torch.float32)
    primals_7 = rand_strided((1, ), (1, ), device='cuda:0', dtype=torch.float32)
    primals_8 = rand_strided((1, 4), (4, 1), device='cuda:0', dtype=torch.float32)
    primals_9 = rand_strided((1, ), (1, ), device='cuda:0', dtype=torch.float32)
    fn = lambda: call([primals_1, primals_2, primals_3, primals_4, primals_5, primals_6, primals_7, primals_8, primals_9])
    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)

from torch.nn import Module
import torch
from torch import nn
from typing import Tuple
import torch.utils.data
import torch.nn.functional
import torch.autograd


class ParityPonderGRU(Module):
    """
    ## PonderNet with GRU for Parity Task

    This is a simple model that uses a [GRU Cell](https://pytorch.org/docs/stable/generated/torch.nn.GRUCell.html)
    as the step function.

    This model is for the [Parity Task](../parity.html) where the input is a vector of `n_elems`.
    Each element of the vector is either `0`, `1` or `-1` and the output is the parity
    - a binary value that is true if the number of `1`s is odd and false otherwise.

    The prediction of the model is the log probability of the parity being $1$.
    """

    def __init__(self, n_elems: 'int', n_hidden: 'int', max_steps: 'int'):
        """
        * `n_elems` is the number of elements in the input vector
        * `n_hidden` is the state vector size of the GRU
        * `max_steps` is the maximum number of steps $N$
        """
        super().__init__()
        self.max_steps = max_steps
        self.n_hidden = n_hidden
        self.gru = nn.GRUCell(n_elems, n_hidden)
        self.output_layer = nn.Linear(n_hidden, 1)
        self.lambda_layer = nn.Linear(n_hidden, 1)
        self.lambda_prob = nn.Sigmoid()
        self.is_halt = False

    def forward(self, x: 'torch.Tensor') ->Tuple[torch.Tensor, torch.Tensor,
        torch.Tensor, torch.Tensor]:
        """
        * `x` is the input of shape `[batch_size, n_elems]`

        This outputs a tuple of four tensors:

        1. $p_1 \\dots p_N$ in a tensor of shape `[N, batch_size]`
        2. $\\hat{y}_1 \\dots \\hat{y}_N$ in a tensor of shape `[N, batch_size]` - the log probabilities of the parity being $1$
        3. $p_m$ of shape `[batch_size]`
        4. $\\hat{y}_m$ of shape `[batch_size]` where the computation was halted at step $m$
        """
        batch_size = x.shape[0]
        h = x.new_zeros((x.shape[0], self.n_hidden))
        h = self.gru(x, h)
        p = []
        y = []
        un_halted_prob = h.new_ones((batch_size,))
        halted = h.new_zeros((batch_size,))
        p_m = h.new_zeros((batch_size,))
        y_m = h.new_zeros((batch_size,))
        for n in range(1, self.max_steps + 1):
            if n == self.max_steps:
                lambda_n = h.new_ones(h.shape[0])
            else:
                lambda_n = self.lambda_prob(self.lambda_layer(h))[:, 0]
            y_n = self.output_layer(h)[:, 0]
            p_n = un_halted_prob * lambda_n
            un_halted_prob = un_halted_prob * (1 - lambda_n)
            halt = torch.bernoulli(lambda_n) * (1 - halted)
            p.append(p_n)
            y.append(y_n)
            p_m = p_m * (1 - halt) + p_n * halt
            y_m = y_m * (1 - halt) + y_n * halt
            halted = halted + halt
            h = self.gru(x, h)
            if self.is_halt and halted.sum() == batch_size:
                break
        return torch.stack(p), torch.stack(y), p_m, y_m


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


def get_init_inputs():
    return [[], {'n_elems': 4, 'n_hidden': 4, 'max_steps': 4}]

import torch
from torch import device
from torch._inductor.select_algorithm import extern_kernels
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.nn import Module
from torch import nn
import torch.utils.data
import torch.nn.functional
import torch.autograd
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor


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


@triton.jit
def triton_poi_fused_stack_1(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl
    .constexpr):
    xnumel = 28
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4
    x0 = xindex % 4
    x2 = xindex
    tmp0 = x1
    tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 1, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.load(in_ptr0 + x0, tmp4 & xmask, eviction_policy='evict_last',
        other=0.0)
    tmp6 = tmp0 >= tmp3
    tmp7 = tl.full([1], 2, tl.int64)
    tmp8 = tmp0 < tmp7
    tmp9 = tmp6 & tmp8
    tmp10 = tl.load(in_ptr0 + x0, tmp9 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp11 = tmp0 >= tmp7
    tmp12 = tl.full([1], 3, tl.int64)
    tmp13 = tmp0 < tmp12
    tmp14 = tmp11 & tmp13
    tmp15 = tl.load(in_ptr0 + x0, tmp14 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp16 = tmp0 >= tmp12
    tmp17 = tl.full([1], 4, tl.int64)
    tmp18 = tmp0 < tmp17
    tmp19 = tmp16 & tmp18
    tmp20 = tl.load(in_ptr0 + x0, tmp19 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp21 = tmp0 >= tmp17
    tmp22 = tl.full([1], 5, tl.int64)
    tmp23 = tmp0 < tmp22
    tmp24 = tmp21 & tmp23
    tmp25 = tl.load(in_ptr1 + x0, tmp24 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp26 = tmp0 >= tmp22
    tmp27 = tl.full([1], 6, tl.int64)
    tmp28 = tmp0 < tmp27
    tmp29 = tmp26 & tmp28
    tmp30 = tl.load(in_ptr1 + x0, tmp29 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp31 = tmp0 >= tmp27
    tl.full([1], 7, tl.int64)
    tmp34 = tl.load(in_ptr1 + x0, tmp31 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp35 = tl.where(tmp29, tmp30, tmp34)
    tmp36 = tl.where(tmp24, tmp25, tmp35)
    tmp37 = tl.where(tmp19, tmp20, tmp36)
    tmp38 = tl.where(tmp14, tmp15, tmp37)
    tmp39 = tl.where(tmp9, tmp10, tmp38)
    tmp40 = tl.where(tmp4, tmp5, tmp39)
    tl.store(out_ptr0 + x2, tmp40, xmask)


@triton.jit
def triton_poi_fused_stack_2(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl
    .constexpr):
    xnumel = 7
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp5 = tl.load(in_ptr0 + 0)
    tmp6 = tl.broadcast_to(tmp5, [XBLOCK])
    tmp23 = tl.load(in_ptr1 + 0)
    tmp24 = tl.broadcast_to(tmp23, [XBLOCK])
    tmp0 = x0
    tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 1, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp7 = tmp0 >= tmp3
    tmp8 = tl.full([1], 2, tl.int64)
    tmp9 = tmp0 < tmp8
    tmp10 = tmp7 & tmp9
    tmp11 = tmp0 >= tmp8
    tmp12 = tl.full([1], 3, tl.int64)
    tmp13 = tmp0 < tmp12
    tmp14 = tmp11 & tmp13
    tmp15 = tmp0 >= tmp12
    tmp16 = tl.full([1], 4, tl.int64)
    tmp17 = tmp0 < tmp16
    tmp18 = tmp15 & tmp17
    tmp19 = tmp0 >= tmp16
    tmp20 = tl.full([1], 5, tl.int64)
    tmp21 = tmp0 < tmp20
    tmp22 = tmp19 & tmp21
    tmp25 = tmp0 >= tmp20
    tmp26 = tl.full([1], 6, tl.int64)
    tmp27 = tmp0 < tmp26
    tmp28 = tmp25 & tmp27
    tl.full([1], 7, tl.int64)
    tmp32 = tl.where(tmp28, tmp24, tmp24)
    tmp33 = tl.where(tmp22, tmp24, tmp32)
    tmp34 = tl.where(tmp18, tmp6, tmp33)
    tmp35 = tl.where(tmp14, tmp6, tmp34)
    tmp36 = tl.where(tmp10, tmp6, tmp35)
    tmp37 = tl.where(tmp4, tmp6, tmp36)
    tl.store(out_ptr0 + x0, tmp37, xmask)


@triton.jit
def triton_poi_fused_stack_3(in_ptr0, in_ptr1, in_ptr2, in_ptr3, out_ptr0,
    xnumel, XBLOCK: tl.constexpr):
    xnumel = 112
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4
    x0 = xindex % 4
    x2 = xindex
    tmp0 = x1
    tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.load(in_ptr0 + (x0 + 4 * x1), tmp4 & xmask, other=0.0)
    tmp6 = tmp0 >= tmp3
    tmp7 = tl.full([1], 8, tl.int64)
    tmp8 = tmp0 < tmp7
    tmp9 = tmp6 & tmp8
    tmp10 = tl.load(in_ptr1 + (x0 + 4 * (-4 + x1)), tmp9 & xmask, other=0.0)
    tmp11 = tmp0 >= tmp7
    tmp12 = tl.full([1], 12, tl.int64)
    tmp13 = tmp0 < tmp12
    tmp14 = tmp11 & tmp13
    tmp15 = tl.load(in_ptr2 + (x0 + 4 * (-8 + x1)), tmp14 & xmask, other=0.0)
    tmp16 = tmp0 >= tmp12
    tmp17 = tl.full([1], 16, tl.int64)
    tmp18 = tmp0 < tmp17
    tmp19 = tmp16 & tmp18
    tmp20 = tl.load(in_ptr3 + (x0 + 4 * (-12 + x1)), tmp19 & xmask, other=0.0)
    tmp21 = tmp0 >= tmp17
    tmp22 = tl.full([1], 20, tl.int64)
    tmp23 = tmp0 < tmp22
    tmp24 = tmp21 & tmp23
    tmp25 = tl.load(in_ptr0 + (x0 + 4 * (-16 + x1)), tmp24 & xmask, other=0.0)
    tmp26 = tmp0 >= tmp22
    tmp27 = tl.full([1], 24, tl.int64)
    tmp28 = tmp0 < tmp27
    tmp29 = tmp26 & tmp28
    tmp30 = tl.load(in_ptr1 + (x0 + 4 * (-20 + x1)), tmp29 & xmask, other=0.0)
    tmp31 = tmp0 >= tmp27
    tl.full([1], 28, tl.int64)
    tmp34 = tl.load(in_ptr2 + (x0 + 4 * (-24 + x1)), tmp31 & xmask, other=0.0)
    tmp35 = tl.where(tmp29, tmp30, tmp34)
    tmp36 = tl.where(tmp24, tmp25, tmp35)
    tmp37 = tl.where(tmp19, tmp20, tmp36)
    tmp38 = tl.where(tmp14, tmp15, tmp37)
    tmp39 = tl.where(tmp9, tmp10, tmp38)
    tmp40 = tl.where(tmp4, tmp5, tmp39)
    tl.store(out_ptr0 + x2, tmp40, xmask)


@triton.jit
def triton_poi_fused_add_bernoulli_mul_new_ones_new_zeros_rsub_sigmoid_4(
    in_out_ptr0, in_out_ptr1, in_ptr0, in_ptr1, in_ptr2, in_ptr3, in_ptr4,
    in_ptr5, out_ptr0, out_ptr1, out_ptr2, out_ptr3, out_ptr4, out_ptr5,
    out_ptr6, 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)
    tmp3 = tl.load(in_ptr1 + 4)
    tmp4 = tl.broadcast_to(tmp3, [XBLOCK])
    tmp5 = tl.load(in_ptr2 + (16 + x0), xmask)
    tmp8 = tl.load(in_ptr1 + 5)
    tmp9 = tl.broadcast_to(tmp8, [XBLOCK])
    tmp10 = tl.load(in_ptr2 + (20 + x0), xmask)
    tmp13 = tl.load(in_ptr1 + 6)
    tmp14 = tl.broadcast_to(tmp13, [XBLOCK])
    tmp15 = tl.load(in_ptr2 + (24 + x0), xmask)
    tmp18 = tl.load(in_ptr3 + x0, xmask)
    tmp20 = tl.load(in_ptr4 + x0, xmask)
    tmp22 = tl.load(in_ptr5 + x0, xmask)
    tmp59 = tl.load(in_ptr1 + 0)
    tmp60 = tl.broadcast_to(tmp59, [XBLOCK])
    tmp61 = tl.load(in_ptr2 + x0, xmask)
    tmp66 = tl.load(in_ptr1 + 1)
    tmp67 = tl.broadcast_to(tmp66, [XBLOCK])
    tmp68 = tl.load(in_ptr2 + (4 + x0), xmask)
    tmp73 = tl.load(in_ptr1 + 2)
    tmp74 = tl.broadcast_to(tmp73, [XBLOCK])
    tmp75 = tl.load(in_ptr2 + (8 + x0), xmask)
    tmp80 = tl.load(in_ptr1 + 3)
    tmp81 = tl.broadcast_to(tmp80, [XBLOCK])
    tmp82 = tl.load(in_ptr2 + (12 + x0), xmask)
    tmp1 = 1.0
    tmp2 = tmp0 < tmp1
    tmp6 = tmp4 + tmp5
    tmp7 = tl.sigmoid(tmp6)
    tmp11 = tmp9 + tmp10
    tmp12 = tl.sigmoid(tmp11)
    tmp16 = tmp14 + tmp15
    tmp17 = tl.sigmoid(tmp16)
    tmp19 = tmp18 < tmp17
    tmp21 = tmp20 < tmp12
    tmp23 = tmp22 < tmp7
    tmp24 = tmp23.to(tl.float32)
    tmp25 = tmp1 - tmp24
    tmp26 = 0.0
    tmp27 = tmp26 * tmp25
    tmp28 = tmp7 * tmp24
    tmp29 = tmp27 + tmp28
    tmp30 = tmp21.to(tl.float32)
    tmp31 = tmp30 * tmp25
    tmp32 = tmp1 - tmp31
    tmp33 = tmp29 * tmp32
    tmp34 = tmp1 - tmp7
    tmp35 = tmp34 * tmp12
    tmp36 = tmp35 * tmp31
    tmp37 = tmp33 + tmp36
    tmp38 = tmp19.to(tl.float32)
    tmp39 = tmp24 + tmp31
    tmp40 = tmp1 - tmp39
    tmp41 = tmp38 * tmp40
    tmp42 = tmp1 - tmp41
    tmp43 = tmp37 * tmp42
    tmp44 = tmp1 - tmp12
    tmp45 = tmp34 * tmp44
    tmp46 = tmp45 * tmp17
    tmp47 = tmp46 * tmp41
    tmp48 = tmp43 + tmp47
    tmp49 = tmp2.to(tl.float32)
    tmp50 = tmp39 + tmp41
    tmp51 = tmp1 - tmp50
    tmp52 = tmp49 * tmp51
    tmp53 = tmp1 - tmp52
    tmp54 = tmp48 * tmp53
    tmp55 = tmp1 - tmp17
    tmp56 = tmp45 * tmp55
    tmp57 = tmp56 * tmp52
    tmp58 = tmp54 + tmp57
    tmp62 = tmp60 + tmp61
    tmp63 = tmp62 * tmp24
    tmp64 = tmp27 + tmp63
    tmp65 = tmp64 * tmp32
    tmp69 = tmp67 + tmp68
    tmp70 = tmp69 * tmp31
    tmp71 = tmp65 + tmp70
    tmp72 = tmp71 * tmp42
    tmp76 = tmp74 + tmp75
    tmp77 = tmp76 * tmp41
    tmp78 = tmp72 + tmp77
    tmp79 = tmp78 * tmp53
    tmp83 = tmp81 + tmp82
    tmp84 = tmp83 * tmp52
    tmp85 = tmp79 + tmp84
    tl.store(out_ptr0 + x0, tmp2, xmask)
    tl.store(out_ptr1 + x0, tmp7, xmask)
    tl.store(out_ptr2 + x0, tmp12, xmask)
    tl.store(out_ptr3 + x0, tmp17, xmask)
    tl.store(out_ptr4 + x0, tmp19, xmask)
    tl.store(out_ptr5 + x0, tmp21, xmask)
    tl.store(out_ptr6 + x0, tmp23, xmask)
    tl.store(in_out_ptr0 + x0, tmp58, xmask)
    tl.store(in_out_ptr1 + x0, tmp85, xmask)


@triton.jit
def triton_poi_fused_stack_5(in_ptr0, in_ptr1, in_ptr2, out_ptr0, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = x0
    tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.load(in_ptr0 + x0, tmp4 & xmask, eviction_policy='evict_last',
        other=0.0)
    tmp6 = tmp0 >= tmp3
    tmp7 = tl.full([1], 8, tl.int64)
    tmp8 = tmp0 < tmp7
    tmp9 = tmp6 & tmp8
    tmp10 = tl.load(in_ptr0 + (-4 + x0), tmp9 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp11 = 1.0
    tmp12 = tmp11 - tmp10
    tmp13 = tl.load(in_ptr1 + (-4 + x0), tmp9 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp14 = tmp12 * tmp13
    tmp15 = tl.full(tmp14.shape, 0.0, tmp14.dtype)
    tmp16 = tl.where(tmp9, tmp14, tmp15)
    tmp17 = tmp0 >= tmp7
    tmp18 = tl.full([1], 12, tl.int64)
    tmp19 = tmp0 < tmp18
    tmp20 = tmp17 & tmp19
    tmp21 = tl.load(in_ptr0 + (-8 + x0), tmp20 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp22 = tmp11 - tmp21
    tmp23 = tl.load(in_ptr1 + (-8 + x0), tmp20 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp24 = tmp11 - tmp23
    tmp25 = tmp22 * tmp24
    tmp26 = tl.load(in_ptr2 + (-8 + x0), tmp20 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp27 = tmp25 * tmp26
    tmp28 = tl.full(tmp27.shape, 0.0, tmp27.dtype)
    tmp29 = tl.where(tmp20, tmp27, tmp28)
    tmp30 = tmp0 >= tmp18
    tl.full([1], 16, tl.int64)
    tmp33 = tl.load(in_ptr0 + (-12 + x0), tmp30 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp34 = tmp11 - tmp33
    tmp35 = tl.load(in_ptr1 + (-12 + x0), tmp30 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp36 = tmp11 - tmp35
    tmp37 = tmp34 * tmp36
    tmp38 = tl.load(in_ptr2 + (-12 + x0), tmp30 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp39 = tmp11 - tmp38
    tmp40 = tmp37 * tmp39
    tmp41 = tl.full(tmp40.shape, 0.0, tmp40.dtype)
    tmp42 = tl.where(tmp30, tmp40, tmp41)
    tmp43 = tl.where(tmp20, tmp29, tmp42)
    tmp44 = tl.where(tmp9, tmp16, tmp43)
    tmp45 = tl.where(tmp4, tmp5, tmp44)
    tl.store(out_ptr0 + x0, tmp45, xmask)


@triton.jit
def triton_poi_fused_stack_6(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
    x0 = xindex
    tmp5 = tl.load(in_ptr0 + 0)
    tmp6 = tl.broadcast_to(tmp5, [XBLOCK])
    tmp15 = tl.load(in_ptr0 + 1)
    tmp16 = tl.broadcast_to(tmp15, [XBLOCK])
    tmp25 = tl.load(in_ptr0 + 2)
    tmp26 = tl.broadcast_to(tmp25, [XBLOCK])
    tmp34 = tl.load(in_ptr0 + 3)
    tmp35 = tl.broadcast_to(tmp34, [XBLOCK])
    tmp0 = x0
    tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp7 = tl.load(in_ptr1 + x0, tmp4 & xmask, eviction_policy='evict_last',
        other=0.0)
    tmp8 = tmp6 + tmp7
    tmp9 = tl.full(tmp8.shape, 0.0, tmp8.dtype)
    tmp10 = tl.where(tmp4, tmp8, tmp9)
    tmp11 = tmp0 >= tmp3
    tmp12 = tl.full([1], 8, tl.int64)
    tmp13 = tmp0 < tmp12
    tmp14 = tmp11 & tmp13
    tmp17 = tl.load(in_ptr1 + (4 + (-4 + x0)), tmp14 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp18 = tmp16 + tmp17
    tmp19 = tl.full(tmp18.shape, 0.0, tmp18.dtype)
    tmp20 = tl.where(tmp14, tmp18, tmp19)
    tmp21 = tmp0 >= tmp12
    tmp22 = tl.full([1], 12, tl.int64)
    tmp23 = tmp0 < tmp22
    tmp24 = tmp21 & tmp23
    tmp27 = tl.load(in_ptr1 + (8 + (-8 + x0)), tmp24 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp28 = tmp26 + tmp27
    tmp29 = tl.full(tmp28.shape, 0.0, tmp28.dtype)
    tmp30 = tl.where(tmp24, tmp28, tmp29)
    tmp31 = tmp0 >= tmp22
    tl.full([1], 16, tl.int64)
    tmp36 = tl.load(in_ptr1 + (12 + (-12 + x0)), tmp31 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp37 = tmp35 + tmp36
    tmp38 = tl.full(tmp37.shape, 0.0, tmp37.dtype)
    tmp39 = tl.where(tmp31, tmp37, tmp38)
    tmp40 = tl.where(tmp24, tmp30, tmp39)
    tmp41 = tl.where(tmp14, tmp20, tmp40)
    tmp42 = tl.where(tmp4, tmp10, tmp41)
    tl.store(out_ptr0 + x0, tmp42, xmask)


def call(args):
    (primals_1, primals_2, primals_3, primals_4, primals_5, primals_6,
        primals_7, primals_8, primals_9) = args
    args.clear()
    assert_size_stride(primals_1, (4, 4), (4, 1))
    assert_size_stride(primals_2, (12, 4), (4, 1))
    assert_size_stride(primals_3, (12, 4), (4, 1))
    assert_size_stride(primals_4, (12,), (1,))
    assert_size_stride(primals_5, (12,), (1,))
    assert_size_stride(primals_6, (1, 4), (4, 1))
    assert_size_stride(primals_7, (1,), (1,))
    assert_size_stride(primals_8, (1, 4), (4, 1))
    assert_size_stride(primals_9, (1,), (1,))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_new_zeros_0[grid(16)](buf0, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf1 = empty_strided_cuda((4, 12), (12, 1), torch.float32)
        extern_kernels.mm(primals_1, reinterpret_tensor(primals_2, (4, 12),
            (1, 4), 0), out=buf1)
        del primals_2
        buf2 = empty_strided_cuda((4, 12), (12, 1), torch.float32)
        extern_kernels.mm(buf0, reinterpret_tensor(primals_3, (4, 12), (1, 
            4), 0), out=buf2)
        buf3 = torch.ops.aten._thnn_fused_gru_cell.default(buf1, buf2, buf0,
            primals_4, primals_5)
        buf4 = buf3[0]
        buf5 = buf3[1]
        del buf3
        buf6 = empty_strided_cuda((7, 4), (4, 1), torch.float32)
        triton_poi_fused_stack_1[grid(28)](primals_8, primals_6, buf6, 28,
            XBLOCK=32, num_warps=1, num_stages=1)
        del primals_6
        del primals_8
        buf7 = empty_strided_cuda((7,), (1,), torch.float32)
        triton_poi_fused_stack_2[grid(7)](primals_9, primals_7, buf7, 7,
            XBLOCK=8, num_warps=1, num_stages=1)
        del primals_7
        del primals_9
        buf8 = buf2
        del buf2
        extern_kernels.mm(buf4, reinterpret_tensor(primals_3, (4, 12), (1, 
            4), 0), out=buf8)
        buf9 = torch.ops.aten._thnn_fused_gru_cell.default(buf1, buf8, buf4,
            primals_4, primals_5)
        buf10 = buf9[0]
        buf11 = buf9[1]
        del buf9
        buf12 = buf8
        del buf8
        extern_kernels.mm(buf10, reinterpret_tensor(primals_3, (4, 12), (1,
            4), 0), out=buf12)
        buf13 = torch.ops.aten._thnn_fused_gru_cell.default(buf1, buf12,
            buf10, primals_4, primals_5)
        buf14 = buf13[0]
        buf15 = buf13[1]
        del buf13
        buf16 = buf12
        del buf12
        extern_kernels.mm(buf14, reinterpret_tensor(primals_3, (4, 12), (1,
            4), 0), out=buf16)
        buf17 = torch.ops.aten._thnn_fused_gru_cell.default(buf1, buf16,
            buf14, primals_4, primals_5)
        del buf1
        del buf16
        del primals_4
        del primals_5
        buf18 = buf17[0]
        buf19 = buf17[1]
        del buf17
        buf20 = empty_strided_cuda((28, 4), (4, 1), torch.float32)
        triton_poi_fused_stack_3[grid(112)](buf4, buf10, buf14, buf18,
            buf20, 112, XBLOCK=128, num_warps=4, num_stages=1)
        buf21 = empty_strided_cuda((7, 4, 1), (4, 1, 1), torch.float32)
        extern_kernels.bmm(reinterpret_tensor(buf20, (7, 4, 4), (16, 4, 1),
            0), reinterpret_tensor(buf6, (7, 4, 1), (4, 1, 4), 0), out=buf21)
        buf23 = torch.ops.aten.rand.default([4], dtype=torch.float32,
            device=device(type='cuda', index=0), pin_memory=False)
        buf24 = buf23
        del buf23
        buf27 = torch.ops.aten.rand.default([4], dtype=torch.float32,
            device=device(type='cuda', index=0), pin_memory=False)
        buf28 = buf27
        del buf27
        buf31 = torch.ops.aten.rand.default([4], dtype=torch.float32,
            device=device(type='cuda', index=0), pin_memory=False)
        buf32 = buf31
        del buf31
        buf36 = torch.ops.aten.rand.default([4], dtype=torch.float32,
            device=device(type='cuda', index=0), pin_memory=False)
        buf37 = buf36
        del buf36
        buf38 = empty_strided_cuda((4,), (1,), torch.bool)
        buf30 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        buf26 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        buf22 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        buf25 = empty_strided_cuda((4,), (1,), torch.bool)
        buf29 = empty_strided_cuda((4,), (1,), torch.bool)
        buf33 = empty_strided_cuda((4,), (1,), torch.bool)
        buf34 = empty_strided_cuda((4,), (1,), torch.float32)
        buf39 = buf34
        del buf34
        buf35 = empty_strided_cuda((4,), (1,), torch.float32)
        buf40 = buf35
        del buf35
        triton_poi_fused_add_bernoulli_mul_new_ones_new_zeros_rsub_sigmoid_4[
            grid(4)](buf39, buf40, buf37, buf7, buf21, buf24, buf28, buf32,
            buf38, buf30, buf26, buf22, buf25, buf29, buf33, 4, XBLOCK=4,
            num_warps=1, num_stages=1)
        del buf24
        del buf28
        del buf32
        del buf37
        buf41 = reinterpret_tensor(buf18, (16,), (1,), 0)
        del buf18
        triton_poi_fused_stack_5[grid(16)](buf30, buf26, buf22, buf41, 16,
            XBLOCK=16, num_warps=1, num_stages=1)
        buf42 = empty_strided_cuda((16,), (1,), torch.float32)
        triton_poi_fused_stack_6[grid(16)](buf7, buf21, buf42, 16, XBLOCK=
            16, num_warps=1, num_stages=1)
        del buf21
        del buf7
    return (reinterpret_tensor(buf41, (4, 4), (4, 1), 0),
        reinterpret_tensor(buf42, (4, 4), (4, 1), 0), buf39, buf40,
        primals_1, buf0, buf4, buf5, buf10, buf11, buf14, buf15, buf19,
        buf22, buf25, buf26, buf29, buf30, buf33, buf38, reinterpret_tensor
        (buf6, (7, 1, 4), (4, 4, 1), 0), reinterpret_tensor(buf20, (7, 4, 4
        ), (16, 1, 4), 0), primals_3)


class ParityPonderGRUNew(Module):
    """
    ## PonderNet with GRU for Parity Task

    This is a simple model that uses a [GRU Cell](https://pytorch.org/docs/stable/generated/torch.nn.GRUCell.html)
    as the step function.

    This model is for the [Parity Task](../parity.html) where the input is a vector of `n_elems`.
    Each element of the vector is either `0`, `1` or `-1` and the output is the parity
    - a binary value that is true if the number of `1`s is odd and false otherwise.

    The prediction of the model is the log probability of the parity being $1$.
    """

    def __init__(self, n_elems: 'int', n_hidden: 'int', max_steps: 'int'):
        """
        * `n_elems` is the number of elements in the input vector
        * `n_hidden` is the state vector size of the GRU
        * `max_steps` is the maximum number of steps $N$
        """
        super().__init__()
        self.max_steps = max_steps
        self.n_hidden = n_hidden
        self.gru = nn.GRUCell(n_elems, n_hidden)
        self.output_layer = nn.Linear(n_hidden, 1)
        self.lambda_layer = nn.Linear(n_hidden, 1)
        self.lambda_prob = nn.Sigmoid()
        self.is_halt = False

    def forward(self, input_0):
        primals_2 = self.gru.weight_ih
        primals_3 = self.gru.weight_hh
        primals_4 = self.gru.bias_ih
        primals_5 = self.gru.bias_hh
        primals_6 = self.output_layer.weight
        primals_7 = self.output_layer.bias
        primals_8 = self.lambda_layer.weight
        primals_9 = self.lambda_layer.bias
        primals_1 = input_0
        output = call([primals_1, primals_2, primals_3, primals_4,
            primals_5, primals_6, primals_7, primals_8, primals_9])
        return output[0], output[1], output[2], output[3]
