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.triton_helpers import libdevice
import math
from torch import nn
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_mul_sub_0(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr
    ):
    xnumel = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x1 = xindex // 4 % 4
    x2 = xindex // 16
    x3 = xindex
    tmp0 = tl.full([1], 0, tl.int64)
    tmp1 = tl.full([1], 1, tl.int64)
    tmp2 = tmp0 < tmp1
    tmp3 = tl.where(tmp2, tmp0, tmp1)
    tmp4 = tl.load(in_ptr0 + (tmp3 + 4 * x0 + 16 * ((x0 + 4 * x1) // 16) + 
        64 * x2 + 64 * ((x0 + 4 * x1) // 64)), xmask, eviction_policy=
        'evict_last')
    tmp5 = tl.load(in_ptr0 + (tmp3 + 4 * x1 + 64 * x2 + 64 * ((x0 + 4 * x1) //
        64)), xmask, eviction_policy='evict_last')
    tmp6 = tmp4 - tmp5
    tmp7 = tmp6 * tmp6
    tmp8 = tmp1 < tmp1
    tmp9 = tl.where(tmp8, tmp0, tmp1)
    tmp10 = tl.load(in_ptr0 + (tmp9 + 4 * x0 + 16 * ((x0 + 4 * x1) // 16) +
        64 * x2 + 64 * ((x0 + 4 * x1) // 64)), xmask, eviction_policy=
        'evict_last')
    tmp11 = tl.load(in_ptr0 + (tmp9 + 4 * x1 + 64 * x2 + 64 * ((x0 + 4 * x1
        ) // 64)), xmask, eviction_policy='evict_last')
    tmp12 = tmp10 - tmp11
    tmp13 = tmp12 * tmp12
    tmp14 = tmp7 + tmp13
    tmp15 = libdevice.sqrt(tmp14)
    tmp16 = 2.0
    tmp17 = tmp15 * tmp16
    tmp18 = tl.load(in_ptr0 + (16 + tmp3 + 4 * x0 + 16 * ((x0 + 4 * x1) // 
        16) + 64 * x2 + 64 * ((16 + x0 + 4 * x1) // 64)), xmask,
        eviction_policy='evict_last')
    tmp19 = tl.load(in_ptr0 + (16 + tmp3 + 4 * x1 + 16 * ((x0 + 4 * x1) // 
        16) + 64 * x2 + 64 * ((16 + x0 + 4 * x1) // 64)), xmask,
        eviction_policy='evict_last')
    tmp20 = tmp18 - tmp19
    tmp21 = tmp20 * tmp20
    tmp22 = tl.load(in_ptr0 + (16 + tmp9 + 4 * x0 + 16 * ((x0 + 4 * x1) // 
        16) + 64 * x2 + 64 * ((16 + x0 + 4 * x1) // 64)), xmask,
        eviction_policy='evict_last')
    tmp23 = tl.load(in_ptr0 + (16 + tmp9 + 4 * x1 + 16 * ((x0 + 4 * x1) // 
        16) + 64 * x2 + 64 * ((16 + x0 + 4 * x1) // 64)), xmask,
        eviction_policy='evict_last')
    tmp24 = tmp22 - tmp23
    tmp25 = tmp24 * tmp24
    tmp26 = tmp21 + tmp25
    tmp27 = libdevice.sqrt(tmp26)
    tmp28 = tmp17 - tmp27
    tl.store(out_ptr0 + x3, tmp28, xmask)


@triton.jit
def triton_poi_fused_mul_sub_1(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr
    ):
    xnumel = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x1 = xindex // 4 % 4
    x2 = xindex // 16
    x3 = xindex
    tmp0 = tl.full([1], 0, tl.int64)
    tmp1 = tl.full([1], 1, tl.int64)
    tmp2 = tmp0 < tmp1
    tmp3 = tl.where(tmp2, tmp0, tmp1)
    tmp4 = tl.load(in_ptr0 + (48 + tmp3 + 4 * x0 + 16 * ((x0 + 4 * x1) // 
        16) + 64 * x2 + 64 * ((48 + x0 + 4 * x1) // 64)), xmask,
        eviction_policy='evict_last')
    tmp5 = tl.load(in_ptr0 + (48 + tmp3 + 4 * x1 + 16 * ((x0 + 4 * x1) // 
        16) + 64 * x2 + 64 * ((48 + x0 + 4 * x1) // 64)), xmask,
        eviction_policy='evict_last')
    tmp6 = tmp4 - tmp5
    tmp7 = tmp6 * tmp6
    tmp8 = tmp1 < tmp1
    tmp9 = tl.where(tmp8, tmp0, tmp1)
    tmp10 = tl.load(in_ptr0 + (48 + tmp9 + 4 * x0 + 16 * ((x0 + 4 * x1) // 
        16) + 64 * x2 + 64 * ((48 + x0 + 4 * x1) // 64)), xmask,
        eviction_policy='evict_last')
    tmp11 = tl.load(in_ptr0 + (48 + tmp9 + 4 * x1 + 16 * ((x0 + 4 * x1) // 
        16) + 64 * x2 + 64 * ((48 + x0 + 4 * x1) // 64)), xmask,
        eviction_policy='evict_last')
    tmp12 = tmp10 - tmp11
    tmp13 = tmp12 * tmp12
    tmp14 = tmp7 + tmp13
    tmp15 = libdevice.sqrt(tmp14)
    tmp16 = 2.0
    tmp17 = tmp15 * tmp16
    tmp18 = tl.load(in_ptr0 + (32 + tmp3 + 4 * x0 + 16 * ((x0 + 4 * x1) // 
        16) + 64 * x2 + 64 * ((32 + x0 + 4 * x1) // 64)), xmask,
        eviction_policy='evict_last')
    tmp19 = tl.load(in_ptr0 + (32 + tmp3 + 4 * x1 + 16 * ((x0 + 4 * x1) // 
        16) + 64 * x2 + 64 * ((32 + x0 + 4 * x1) // 64)), xmask,
        eviction_policy='evict_last')
    tmp20 = tmp18 - tmp19
    tmp21 = tmp20 * tmp20
    tmp22 = tl.load(in_ptr0 + (32 + tmp9 + 4 * x0 + 16 * ((x0 + 4 * x1) // 
        16) + 64 * x2 + 64 * ((32 + x0 + 4 * x1) // 64)), xmask,
        eviction_policy='evict_last')
    tmp23 = tl.load(in_ptr0 + (32 + tmp9 + 4 * x1 + 16 * ((x0 + 4 * x1) // 
        16) + 64 * x2 + 64 * ((32 + x0 + 4 * x1) // 64)), xmask,
        eviction_policy='evict_last')
    tmp24 = tmp22 - tmp23
    tmp25 = tmp24 * tmp24
    tmp26 = tmp21 + tmp25
    tmp27 = libdevice.sqrt(tmp26)
    tmp28 = tmp17 - tmp27
    tl.store(out_ptr0 + x3, tmp28, xmask)


@triton.jit
def triton_poi_fused_cat_2(in_ptr0, in_ptr1, in_ptr2, out_ptr0, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 384
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex // 16 % 6
    x3 = xindex // 96
    x4 = xindex % 16
    x0 = xindex % 4
    x1 = xindex // 4 % 4
    x5 = xindex
    tmp0 = x2
    tmp1 = tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 1, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.load(in_ptr0 + (x4 + 16 * x3), tmp4 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp6 = tmp0 >= tmp3
    tmp7 = tl.full([1], 5, tl.int64)
    tmp8 = tmp0 < tmp7
    tmp9 = tmp6 & tmp8
    tmp10 = tmp1 < tmp3
    tmp11 = tl.where(tmp10, tmp1, tmp3)
    tmp12 = tl.load(in_ptr1 + (tmp11 + 4 * x0 + 16 * ((x0 + 4 * x1 + 16 * (
        -1 + x2)) // 16 % 4) + 64 * ((x0 + 4 * x1 + 16 * (-1 + x2) + 64 *
        x3) // 64 % 4)), tmp9 & xmask, eviction_policy='evict_last', other=0.0)
    tmp13 = tl.load(in_ptr1 + (tmp11 + 4 * x1 + 16 * ((x0 + 4 * x1 + 16 * (
        -1 + x2)) // 16 % 4) + 64 * ((x0 + 4 * x1 + 16 * (-1 + x2) + 64 *
        x3) // 64 % 4)), tmp9 & xmask, eviction_policy='evict_last', other=0.0)
    tmp14 = tmp12 - tmp13
    tmp15 = tmp14 * tmp14
    tmp16 = tmp3 < tmp3
    tmp17 = tl.where(tmp16, tmp1, tmp3)
    tmp18 = tl.load(in_ptr1 + (tmp17 + 4 * x0 + 16 * ((x0 + 4 * x1 + 16 * (
        -1 + x2)) // 16 % 4) + 64 * ((x0 + 4 * x1 + 16 * (-1 + x2) + 64 *
        x3) // 64 % 4)), tmp9 & xmask, eviction_policy='evict_last', other=0.0)
    tmp19 = tl.load(in_ptr1 + (tmp17 + 4 * x1 + 16 * ((x0 + 4 * x1 + 16 * (
        -1 + x2)) // 16 % 4) + 64 * ((x0 + 4 * x1 + 16 * (-1 + x2) + 64 *
        x3) // 64 % 4)), tmp9 & xmask, eviction_policy='evict_last', other=0.0)
    tmp20 = tmp18 - tmp19
    tmp21 = tmp20 * tmp20
    tmp22 = tmp15 + tmp21
    tmp23 = libdevice.sqrt(tmp22)
    tmp24 = tl.full(tmp23.shape, 0.0, tmp23.dtype)
    tmp25 = tl.where(tmp9, tmp23, tmp24)
    tmp26 = tmp0 >= tmp7
    tl.full([1], 6, tl.int64)
    tmp29 = tl.load(in_ptr2 + (x4 + 16 * x3), tmp26 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp30 = tl.where(tmp9, tmp25, tmp29)
    tmp31 = tl.where(tmp4, tmp5, tmp30)
    tl.store(out_ptr0 + x5, tmp31, xmask)


@triton.jit
def triton_poi_fused_add_div_sub_3(in_ptr0, in_ptr1, out_ptr0, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x1 = xindex // 4 % 4
    x2 = xindex // 16 % 4
    x3 = xindex // 64
    x4 = xindex
    tmp16 = tl.load(in_ptr1 + 0)
    tmp17 = tl.broadcast_to(tmp16, [XBLOCK])
    tmp19 = tl.load(in_ptr1 + 1)
    tmp20 = tl.broadcast_to(tmp19, [XBLOCK])
    tmp0 = tl.full([1], 0, tl.int64)
    tmp1 = tl.full([1], 1, tl.int64)
    tmp2 = tmp0 < tmp1
    tmp3 = tl.where(tmp2, tmp0, tmp1)
    tmp4 = tl.load(in_ptr0 + (tmp3 + 4 * x0 + 16 * x2 + 16 * ((x0 + 4 * x1) //
        16) + 64 * x3 + 64 * ((x0 + 4 * x1 + 16 * x2) // 64)), xmask,
        eviction_policy='evict_last')
    tmp5 = tl.load(in_ptr0 + (tmp3 + 4 * x1 + 16 * x2 + 16 * ((x0 + 4 * x1) //
        16) + 64 * x3 + 64 * ((x0 + 4 * x1 + 16 * x2) // 64)), xmask,
        eviction_policy='evict_last')
    tmp6 = tmp4 - tmp5
    tmp7 = tmp6 * tmp6
    tmp8 = tmp1 < tmp1
    tmp9 = tl.where(tmp8, tmp0, tmp1)
    tmp10 = tl.load(in_ptr0 + (tmp9 + 4 * x0 + 16 * x2 + 16 * ((x0 + 4 * x1
        ) // 16) + 64 * x3 + 64 * ((x0 + 4 * x1 + 16 * x2) // 64)), xmask,
        eviction_policy='evict_last')
    tmp11 = tl.load(in_ptr0 + (tmp9 + 4 * x1 + 16 * x2 + 16 * ((x0 + 4 * x1
        ) // 16) + 64 * x3 + 64 * ((x0 + 4 * x1 + 16 * x2) // 64)), xmask,
        eviction_policy='evict_last')
    tmp12 = tmp10 - tmp11
    tmp13 = tmp12 * tmp12
    tmp14 = tmp7 + tmp13
    tmp15 = libdevice.sqrt(tmp14)
    tmp18 = tmp15 - tmp17
    tmp21 = tmp20 - tmp17
    tmp22 = 1e-05
    tmp23 = tmp21 + tmp22
    tmp24 = tmp18 / tmp23
    tl.store(out_ptr0 + x4, tmp24, xmask)


@triton.jit
def triton_poi_fused_div_sub_4(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK:
    tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 64
    x1 = xindex // 64
    x2 = xindex
    tmp0 = tl.load(in_ptr0 + (32 + x0 + 96 * x1), xmask)
    tmp1 = tl.load(in_ptr0 + (x0 + 96 * x1), xmask)
    tmp5 = tl.load(in_ptr1 + 0)
    tmp6 = tl.broadcast_to(tmp5, [XBLOCK])
    tmp8 = tl.load(in_ptr1 + 1)
    tmp9 = tl.broadcast_to(tmp8, [XBLOCK])
    tmp2 = tmp0 - tmp1
    tmp3 = 0.5
    tmp4 = tmp2 * tmp3
    tmp7 = tmp4 - tmp6
    tmp10 = tmp7 / tmp9
    tl.store(out_ptr0 + x2, tmp10, xmask)


@triton.jit
def triton_poi_fused_div_sigmoid_sigmoid_backward_sub_5(in_ptr0, in_ptr1,
    in_ptr2, out_ptr0, out_ptr1, out_ptr2, out_ptr3, xnumel, XBLOCK: tl.
    constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x1 = xindex // 4 % 4
    x2 = xindex // 16 % 4
    x3 = xindex // 64
    x5 = xindex // 16
    x6 = xindex // 4
    x4 = xindex
    tmp0 = tl.load(in_ptr0 + (4 * x0 + 16 * x2 + 16 * ((x0 + 4 * x1) // 16) +
        64 * x3 + 64 * ((x0 + 4 * x1 + 16 * x2) // 64)), xmask,
        eviction_policy='evict_last')
    tmp1 = tl.load(in_ptr0 + (4 * x1 + 16 * x2 + 16 * ((x0 + 4 * x1) // 16) +
        64 * x3 + 64 * ((x0 + 4 * x1 + 16 * x2) // 64)), xmask,
        eviction_policy='evict_last')
    tmp5 = tl.load(in_ptr0 + (2 + 4 * x0 + 16 * x5), xmask, eviction_policy
        ='evict_last')
    tmp6 = tl.load(in_ptr0 + (2 + 4 * x6), xmask, eviction_policy='evict_last')
    tmp11 = tl.load(in_ptr1 + 0)
    tmp12 = tl.broadcast_to(tmp11, [XBLOCK])
    tmp20 = tl.load(in_ptr2 + 0)
    tmp21 = tl.broadcast_to(tmp20, [XBLOCK])
    tmp2 = tmp0 - tmp1
    tmp3 = tmp2 * tmp2
    tmp4 = libdevice.sqrt(tmp3)
    tmp7 = tmp5 + tmp6
    tmp8 = 0.5
    tmp9 = tmp7 * tmp8
    tmp10 = tmp4 - tmp9
    tmp13 = tmp10 - tmp12
    tmp14 = 1.0
    tmp15 = tmp13 * tmp14
    tmp16 = tl.sigmoid(tmp15)
    tmp17 = tmp5 - tmp6
    tmp18 = tmp17 * tmp8
    tmp19 = tmp4 - tmp18
    tmp22 = tmp19 - tmp21
    tmp23 = tmp22 * tmp14
    tmp24 = tl.sigmoid(tmp23)
    tmp25 = tmp14 - tmp24
    tmp26 = tmp24 * tmp25
    tmp27 = tmp14 - tmp16
    tmp28 = tmp16 * tmp27
    tl.store(out_ptr0 + 24 * x4, tmp16, xmask)
    tl.store(out_ptr1 + 24 * x4, tmp24, xmask)
    tl.store(out_ptr2 + x4, tmp26, xmask)
    tl.store(out_ptr3 + x4, tmp28, xmask)


@triton.jit
def triton_poi_fused_div_sigmoid_sigmoid_backward_sub_6(in_ptr0, in_ptr1,
    in_ptr2, out_ptr0, out_ptr1, out_ptr2, out_ptr3, xnumel, XBLOCK: tl.
    constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x1 = xindex // 4 % 4
    x2 = xindex // 16 % 4
    x3 = xindex // 64
    x5 = xindex // 16
    x6 = xindex // 4
    x4 = xindex
    tmp0 = tl.load(in_ptr0 + (1 + 4 * x0 + 16 * x2 + 16 * ((x0 + 4 * x1) //
        16) + 64 * x3 + 64 * ((x0 + 4 * x1 + 16 * x2) // 64)), xmask,
        eviction_policy='evict_last')
    tmp1 = tl.load(in_ptr0 + (1 + 4 * x1 + 16 * x2 + 16 * ((x0 + 4 * x1) //
        16) + 64 * x3 + 64 * ((x0 + 4 * x1 + 16 * x2) // 64)), xmask,
        eviction_policy='evict_last')
    tmp5 = tl.load(in_ptr0 + (3 + 4 * x0 + 16 * x5), xmask, eviction_policy
        ='evict_last')
    tmp6 = tl.load(in_ptr0 + (3 + 4 * x6), xmask, eviction_policy='evict_last')
    tmp11 = tl.load(in_ptr1 + 0)
    tmp12 = tl.broadcast_to(tmp11, [XBLOCK])
    tmp20 = tl.load(in_ptr2 + 0)
    tmp21 = tl.broadcast_to(tmp20, [XBLOCK])
    tmp2 = tmp0 - tmp1
    tmp3 = tmp2 * tmp2
    tmp4 = libdevice.sqrt(tmp3)
    tmp7 = tmp5 + tmp6
    tmp8 = 0.5
    tmp9 = tmp7 * tmp8
    tmp10 = tmp4 - tmp9
    tmp13 = tmp10 - tmp12
    tmp14 = 1.0
    tmp15 = tmp13 * tmp14
    tmp16 = tl.sigmoid(tmp15)
    tmp17 = tmp5 - tmp6
    tmp18 = tmp17 * tmp8
    tmp19 = tmp4 - tmp18
    tmp22 = tmp19 - tmp21
    tmp23 = tmp22 * tmp14
    tmp24 = tl.sigmoid(tmp23)
    tmp25 = tmp14 - tmp24
    tmp26 = tmp24 * tmp25
    tmp27 = tmp14 - tmp16
    tmp28 = tmp16 * tmp27
    tl.store(out_ptr0 + 24 * x4, tmp16, xmask)
    tl.store(out_ptr1 + 24 * x4, tmp24, xmask)
    tl.store(out_ptr2 + x4, tmp26, xmask)
    tl.store(out_ptr3 + x4, tmp28, xmask)


@triton.jit
def triton_poi_fused_div_sigmoid_sub_7(in_ptr0, in_ptr1, out_ptr0, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 2560
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 10
    x0 = xindex % 10
    tmp0 = tl.load(in_ptr0 + x1, xmask, eviction_policy='evict_last')
    tmp1 = tl.load(in_ptr1 + x0, xmask, eviction_policy='evict_last')
    tmp2 = tmp0 - tmp1
    tmp3 = 1.0
    tmp4 = tmp2 * tmp3
    tmp5 = tl.sigmoid(tmp4)
    tl.store(out_ptr0 + (x0 + 24 * x1), tmp5, xmask)


@triton.jit
def triton_poi_fused_div_sigmoid_sub_8(in_ptr0, in_ptr1, out_ptr0, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 2560
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 10
    x0 = xindex % 10
    tmp0 = tl.load(in_ptr0 + x1, xmask, eviction_policy='evict_last')
    tmp1 = tl.load(in_ptr1 + x0, xmask, eviction_policy='evict_last')
    tmp2 = tmp0 - tmp1
    tmp3 = 1.0
    tmp4 = tmp2 * tmp3
    tmp5 = tl.sigmoid(tmp4)
    tl.store(out_ptr0 + (x0 + 24 * x1), tmp5, 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, 4), (64, 16, 4, 1))
    assert_size_stride(primals_2, (2,), (1,))
    assert_size_stride(primals_3, (10,), (1,))
    assert_size_stride(primals_4, (2,), (1,))
    assert_size_stride(primals_5, (10,), (1,))
    assert_size_stride(primals_6, (1,), (1,))
    assert_size_stride(primals_7, (1,), (1,))
    assert_size_stride(primals_8, (1,), (1,))
    assert_size_stride(primals_9, (1,), (1,))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 1, 4, 4, 1), (16, 64, 4, 1, 64),
            torch.float32)
        get_raw_stream(0)
        triton_poi_fused_mul_sub_0[grid(64)](primals_1, buf0, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf1 = empty_strided_cuda((4, 1, 4, 4, 1), (16, 64, 4, 1, 64),
            torch.float32)
        triton_poi_fused_mul_sub_1[grid(64)](primals_1, buf1, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf2 = empty_strided_cuda((4, 6, 4, 4, 1), (96, 16, 4, 1, 384),
            torch.float32)
        triton_poi_fused_cat_2[grid(384)](buf0, primals_1, buf1, buf2, 384,
            XBLOCK=256, num_warps=4, num_stages=1)
        del buf0
        del buf1
        buf3 = empty_strided_cuda((4, 4, 4, 4, 1), (64, 16, 4, 1, 1), torch
            .float32)
        triton_poi_fused_add_div_sub_3[grid(256)](primals_1, primals_2,
            buf3, 256, XBLOCK=256, num_warps=4, num_stages=1)
        del primals_2
        buf4 = empty_strided_cuda((4, 4, 4, 4, 1), (64, 16, 4, 1, 1), torch
            .float32)
        triton_poi_fused_div_sub_4[grid(256)](buf2, primals_4, buf4, 256,
            XBLOCK=256, num_warps=4, num_stages=1)
        del buf2
        del primals_4
        buf11 = empty_strided_cuda((4, 4, 4, 4, 24), (1536, 384, 96, 24, 1),
            torch.float32)
        buf5 = reinterpret_tensor(buf11, (4, 4, 4, 4, 1), (1536, 384, 96, 
            24, 1), 20)
        buf7 = reinterpret_tensor(buf11, (4, 4, 4, 4, 1), (1536, 384, 96, 
            24, 1), 22)
        buf13 = empty_strided_cuda((4, 4, 4, 4, 1), (64, 16, 4, 1, 1),
            torch.float32)
        buf15 = empty_strided_cuda((4, 4, 4, 4, 1), (64, 16, 4, 1, 1),
            torch.float32)
        triton_poi_fused_div_sigmoid_sigmoid_backward_sub_5[grid(256)](
            primals_1, primals_6, primals_8, buf5, buf7, buf13, buf15, 256,
            XBLOCK=128, num_warps=4, num_stages=1)
        del primals_6
        del primals_8
        buf6 = reinterpret_tensor(buf11, (4, 4, 4, 4, 1), (1536, 384, 96, 
            24, 1), 21)
        buf8 = reinterpret_tensor(buf11, (4, 4, 4, 4, 1), (1536, 384, 96, 
            24, 1), 23)
        buf12 = empty_strided_cuda((4, 4, 4, 4, 1), (64, 16, 4, 1, 1),
            torch.float32)
        buf14 = empty_strided_cuda((4, 4, 4, 4, 1), (64, 16, 4, 1, 1),
            torch.float32)
        triton_poi_fused_div_sigmoid_sigmoid_backward_sub_6[grid(256)](
            primals_1, primals_7, primals_9, buf6, buf8, buf12, buf14, 256,
            XBLOCK=128, num_warps=4, num_stages=1)
        del primals_1
        del primals_7
        del primals_9
        buf9 = reinterpret_tensor(buf11, (4, 4, 4, 4, 10), (1536, 384, 96, 
            24, 1), 0)
        triton_poi_fused_div_sigmoid_sub_7[grid(2560)](buf3, primals_3,
            buf9, 2560, XBLOCK=256, num_warps=4, num_stages=1)
        buf10 = reinterpret_tensor(buf11, (4, 4, 4, 4, 10), (1536, 384, 96,
            24, 1), 10)
        triton_poi_fused_div_sigmoid_sub_8[grid(2560)](buf4, primals_5,
            buf10, 2560, XBLOCK=256, num_warps=4, num_stages=1)
    return buf11, primals_3, primals_5, buf3, buf4, buf12, buf13, buf14, buf15


def apply_last_dim(model, x):
    size = list(x.size())
    y = model(x.contiguous().view(-1, size[-1]))
    size[-1] = y.size(-1)
    y = y.view(torch.Size(size))
    return y


def get_int_dim_index(name):
    if isinstance(name, int):
        return name
    name_list = 'axyz'
    assert name in name_list
    return [i for i in range(len(name_list)) if name_list[i] == name][0] - 1


class Length(nn.Module):

    def __init__(self, dim_index=-1):
        super().__init__()
        self.dim_index = dim_index

    def forward(self, states, dim_index=None):
        if dim_index is None:
            dim_index = self.dim_index
        if isinstance(dim_index, int):
            dim_index = [dim_index]
        else:
            dim_index = [get_int_dim_index(x) for x in dim_index]
        if -1 in dim_index:

            def extractor(x):
                return torch.sqrt(torch.sum(x * x, dim=1, keepdim=True))
        else:

            def extractor(x):
                return torch.sqrt(torch.sum(x[:, dim_index].pow(2), dim=1,
                    keepdim=True))
        return apply_last_dim(extractor, states)

    def show(self, name='Length', indent=0, log=print, **kwargs):
        log(' ' * indent + "- %s(x) = |x's dim %s|" % (name, str(self.
            dim_index)))


class Distance(nn.Module):

    def __init__(self, dim_index=-1):
        super().__init__()
        self.dim_index = dim_index
        self.length = Length(dim_index)

    def forward(self, states1, states2, dim_index=None):
        return self.length(states1 - states2, dim_index)

    def show(self, name='Distance', indent=0, log=print, **kwargs):
        log(' ' * indent + '- %s(x1, x2) = |x1 - x2|' % name)


class Normalize(nn.Module):

    def __init__(self, distribution=None, **kwargs):
        super().__init__()
        self.distribution = distribution
        self.data_ = []
        if distribution is None:
            pass
        elif distribution == 'normal':
            mean = kwargs['mean'] if 'mean' in kwargs else 0
            std = kwargs['std'] if 'std' in kwargs else 1
            self.param = nn.Parameter(torch.Tensor([mean, std]), False)
        elif distribution == 'uniform':
            vmin = kwargs['minv'] if 'minv' in kwargs else 0
            vmax = kwargs['maxv'] if 'maxv' in kwargs else 1
            self.param = nn.Parameter(torch.Tensor([vmin, vmax]), False)
        else:
            raise NotImplementedError()

    def forward(self, x, keep_data=False):
        if keep_data:
            self.data_.append(x.detach().cpu().view(-1))
            return x
        if self.distribution is None:
            return x
        elif self.distribution == 'normal':
            mean = self.param[0]
            std = self.param[1]
            return (x - mean) / std
        elif self.distribution == 'uniform':
            vmin = self.param[0]
            vmax = self.param[1]
            return (x - vmin) / (vmax - vmin + 1e-05)
        else:
            raise NotImplementedError()

    def reset_parameters(self, name=None):
        assert len(self.data_) > 0
        data = torch.cat(self.data_)
        self.data_ = []
        if self.distribution is None:
            pass
        elif self.distribution == 'normal':
            with torch.no_grad():
                self.param[0] = data.mean().item()
                self.param[1] = data.std().item()
            if name is not None:
                None
        elif self.distribution == 'uniform':
            with torch.no_grad():
                self.param[0] = data.min().item()
                self.param[1] = data.max().item()
            if name is not None:
                None
        else:
            raise NotImplementedError()

    def recover_threshold(self, x):
        if self.distribution is None:
            return x
        elif self.distribution == 'normal':
            return x * float(self.param[1]) + float(self.param[0])
        elif self.distribution == 'uniform':
            return x * float(self.param[1] - self.param[0] + 1e-05) + float(
                self.param[0])
        else:
            raise NotImplementedError()

    def init_thresholds(self, x):
        if self.distribution is None:
            nn.init.normal_(x, 0, 1)
        elif self.distribution == 'normal':
            nn.init.normal_(x, 0, 1)
        elif self.distribution == 'uniform':
            nn.init.uniform_(x, 0, 1)
        else:
            raise NotImplementedError()


class SoftCmp(nn.Module):
    """
    Sigmoid((x - y) / e^beta)
    """

    def __init__(self):
        super().__init__()
        self.sigmoid = nn.Sigmoid()

    def forward(self, x, y, beta):
        return self.sigmoid((x - y) / math.exp(beta))


class Inequality(nn.Module):

    def __init__(self, out_dim=1, distribution=None, **kwargs):
        super().__init__()
        self.out_dim = out_dim
        self.thresholds = nn.Parameter(torch.zeros(out_dim), requires_grad=True
            )
        self.distribution = distribution
        self.normalize = Normalize(distribution)
        self.cmp = SoftCmp()
        self.normalize.init_thresholds(self.thresholds)

    def forward(self, states, beta=0, **kwargs):
        """
        :param states: [batch, length, n_agents, ... ]
        """
        states_expand = states.view(*(states.size() + (1,)))
        estimate_parameters = 'estimate_parameters' in kwargs and kwargs[
            'estimate_parameters']
        states_expand = self.normalize(states_expand, keep_data=
            estimate_parameters)
        return self.cmp(states_expand, self.thresholds.view(*([1] * len(
            states.size()) + [self.out_dim])), beta)

    def reset_parameters(self, parameter_name, name=None):
        if parameter_name == 'primitive_inequality':
            self.normalize.reset_parameters(name=name)
            self.normalize.init_thresholds(self.thresholds)

    def get_descriptions(self, name='Inequality'):
        theta = self.thresholds.detach().cpu().view(self.out_dim)
        descroptions = []
        for k in range(theta.size(0)):
            t = self.normalize.recover_threshold(theta[k])
            if 'speed' in name:
                t = t * 8
            if 'acc' in name:
                t = t * 64
            descroptions.append('%s > %.2lf' % (name, t))
        return descroptions


class AlignDifferential(nn.Module):

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

    def new_length(self, length):
        return length

    def forward(self, states):
        """
        :param states: [batch, length, *]
        """
        padded_states = torch.cat([states[:, 0:1] * 2 - states[:, 1:2],
            states, states[:, -1:] * 2 - states[:, -2:-1]], dim=1)
        return (padded_states[:, 2:] - padded_states[:, :-2]) / 2

    def show(self, name='AlignDifferential', indent=0, log=print, **kwargs):
        log(' ' * indent + '- %s(x) = AlignDifferential()' % (name,))


class N_aryPrimitivesSomethingElse(nn.Module):

    def __init__(self):
        super().__init__()
        self.out_dim = 0
        self.ineqs = nn.ModuleDict({})

    def reset_parameters(self, parameter_name):
        for k in self.ineqs:
            self.ineqs[k].reset_parameters(parameter_name, name=k)


class BinaryPrimitivesSomethingElseNew(N_aryPrimitivesSomethingElse):

    def __init__(self, cmp_dim=10):
        super().__init__()
        self.distance = Distance()
        self.differential = AlignDifferential()
        self.ineqs.update({'dist': Inequality(out_dim=cmp_dim, distribution
            ='uniform'), 'app_vel': Inequality(out_dim=cmp_dim,
            distribution='normal'), 'overlap_x': Inequality(out_dim=1,
            distribution=None), 'overlap_y': Inequality(out_dim=1,
            distribution=None), 'contain_x': Inequality(out_dim=1,
            distribution=None), 'contain_y': Inequality(out_dim=1,
            distribution=None)})
        self.out_dim = sum([self.ineqs[k].out_dim for k in self.ineqs])

    def forward(self, input_0):
        primals_3 = self.ineqs.dist.thresholds
        primals_2 = self.ineqs.dist.normalize.param
        primals_5 = self.ineqs.app_vel.thresholds
        primals_4 = self.ineqs.app_vel.normalize.param
        primals_6 = self.ineqs.overlap_x.thresholds
        primals_7 = self.ineqs.overlap_y.thresholds
        primals_8 = self.ineqs.contain_x.thresholds
        primals_9 = self.ineqs.contain_y.thresholds
        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]


Measured latency_ms: 0.015779