import torch
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._inductor.runtime import triton_helpers
from torch._inductor.runtime.triton_helpers import math as tl_math
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data
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_relu_threshold_backward_0(in_ptr0, out_ptr0, out_ptr1,
    xnumel, XBLOCK: tl.constexpr):
    xnumel = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + x0, xmask)
    tmp1 = tl.full([1], 0, tl.int32)
    tmp2 = triton_helpers.maximum(tmp1, tmp0)
    tmp3 = 0.0
    tmp4 = tmp2 <= tmp3
    tl.store(out_ptr0 + x0, tmp2, xmask)
    tl.store(out_ptr1 + x0, tmp4, xmask)


@triton.jit
def triton_poi_fused_mul_1(in_out_ptr0, in_ptr0, in_ptr1, xnumel, XBLOCK:
    tl.constexpr):
    xnumel = 192
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x0 = xindex % 12
    x1 = xindex // 12
    tmp0 = tl.load(in_out_ptr0 + x2, xmask)
    tmp1 = tl.load(in_ptr0 + x0, xmask, eviction_policy='evict_last')
    tmp3 = tl.load(in_ptr1 + x1, xmask, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tmp4 = tmp2 * tmp3
    tl.store(in_out_ptr0 + x2, tmp4, xmask)


@triton.jit
def triton_poi_fused_bmm_2(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (4 + 12 * x0), xmask, eviction_policy='evict_last'
        )
    tl.store(out_ptr0 + x0, tmp0, xmask)


@triton.jit
def triton_poi_fused_bmm_3(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + 12 * x0, xmask, eviction_policy='evict_last')
    tl.store(out_ptr0 + x0, tmp0, xmask)


@triton.jit
def triton_poi_fused_bmm_4(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (5 + 12 * x0), xmask, eviction_policy='evict_last'
        )
    tl.store(out_ptr0 + x0, tmp0, xmask)


@triton.jit
def triton_poi_fused_bmm_5(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (1 + 12 * x0), xmask, eviction_policy='evict_last'
        )
    tl.store(out_ptr0 + x0, tmp0, xmask)


@triton.jit
def triton_poi_fused_bmm_6(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (6 + 12 * x0), xmask, eviction_policy='evict_last'
        )
    tl.store(out_ptr0 + x0, tmp0, xmask)


@triton.jit
def triton_poi_fused_bmm_7(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (2 + 12 * x0), xmask, eviction_policy='evict_last'
        )
    tl.store(out_ptr0 + x0, tmp0, xmask)


@triton.jit
def triton_poi_fused_bmm_8(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (7 + 12 * x0), xmask, eviction_policy='evict_last'
        )
    tl.store(out_ptr0 + x0, tmp0, xmask)


@triton.jit
def triton_poi_fused_bmm_9(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (3 + 12 * x0), xmask, eviction_policy='evict_last'
        )
    tl.store(out_ptr0 + x0, tmp0, xmask)


@triton.jit
def triton_poi_fused__softmax_eq_masked_fill_10(in_ptr0, in_ptr1, in_ptr2,
    in_ptr3, in_ptr4, out_ptr0, out_ptr1, out_ptr2, out_ptr3, out_ptr4,
    out_ptr5, out_ptr6, out_ptr7, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4
    x2 = xindex
    tmp0 = tl.load(in_ptr0 + 4 * x1, xmask, eviction_policy='evict_last')
    tmp3 = tl.load(in_ptr1 + 4 * x2, xmask, eviction_policy='evict_last')
    tmp8 = tl.load(in_ptr0 + (1 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp10 = tl.load(in_ptr1 + (1 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp14 = tl.load(in_ptr0 + (2 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp16 = tl.load(in_ptr1 + (2 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp20 = tl.load(in_ptr0 + (3 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp22 = tl.load(in_ptr1 + (3 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp41 = tl.load(in_ptr2 + 4 * x2, xmask, eviction_policy='evict_last')
    tmp44 = tl.load(in_ptr2 + (1 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp48 = tl.load(in_ptr2 + (2 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp52 = tl.load(in_ptr2 + (3 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp71 = tl.load(in_ptr3 + 4 * x2, xmask, eviction_policy='evict_last')
    tmp74 = tl.load(in_ptr3 + (1 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp78 = tl.load(in_ptr3 + (2 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp82 = tl.load(in_ptr3 + (3 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp101 = tl.load(in_ptr4 + 4 * x2, xmask, eviction_policy='evict_last')
    tmp104 = tl.load(in_ptr4 + (1 + 4 * x2), xmask, eviction_policy=
        'evict_last')
    tmp108 = tl.load(in_ptr4 + (2 + 4 * x2), xmask, eviction_policy=
        'evict_last')
    tmp112 = tl.load(in_ptr4 + (3 + 4 * x2), xmask, eviction_policy=
        'evict_last')
    tmp1 = 0.0
    tmp2 = tmp0 == tmp1
    tmp4 = -1000000000.0
    tmp5 = tl.where(tmp2, tmp4, tmp3)
    tmp6 = 1.0
    tmp7 = tmp5 * tmp6
    tmp9 = tmp8 == tmp1
    tmp11 = tl.where(tmp9, tmp4, tmp10)
    tmp12 = tmp11 * tmp6
    tmp13 = triton_helpers.maximum(tmp7, tmp12)
    tmp15 = tmp14 == tmp1
    tmp17 = tl.where(tmp15, tmp4, tmp16)
    tmp18 = tmp17 * tmp6
    tmp19 = triton_helpers.maximum(tmp13, tmp18)
    tmp21 = tmp20 == tmp1
    tmp23 = tl.where(tmp21, tmp4, tmp22)
    tmp24 = tmp23 * tmp6
    tmp25 = triton_helpers.maximum(tmp19, tmp24)
    tmp26 = tmp7 - tmp25
    tmp27 = tmp26 * tmp6
    tmp28 = tl_math.exp(tmp27)
    tmp29 = tmp12 - tmp25
    tmp30 = tmp29 * tmp6
    tmp31 = tl_math.exp(tmp30)
    tmp32 = tmp28 + tmp31
    tmp33 = tmp18 - tmp25
    tmp34 = tmp33 * tmp6
    tmp35 = tl_math.exp(tmp34)
    tmp36 = tmp32 + tmp35
    tmp37 = tmp24 - tmp25
    tmp38 = tmp37 * tmp6
    tmp39 = tl_math.exp(tmp38)
    tmp40 = tmp36 + tmp39
    tmp42 = tl.where(tmp2, tmp4, tmp41)
    tmp43 = tmp42 * tmp6
    tmp45 = tl.where(tmp9, tmp4, tmp44)
    tmp46 = tmp45 * tmp6
    tmp47 = triton_helpers.maximum(tmp43, tmp46)
    tmp49 = tl.where(tmp15, tmp4, tmp48)
    tmp50 = tmp49 * tmp6
    tmp51 = triton_helpers.maximum(tmp47, tmp50)
    tmp53 = tl.where(tmp21, tmp4, tmp52)
    tmp54 = tmp53 * tmp6
    tmp55 = triton_helpers.maximum(tmp51, tmp54)
    tmp56 = tmp43 - tmp55
    tmp57 = tmp56 * tmp6
    tmp58 = tl_math.exp(tmp57)
    tmp59 = tmp46 - tmp55
    tmp60 = tmp59 * tmp6
    tmp61 = tl_math.exp(tmp60)
    tmp62 = tmp58 + tmp61
    tmp63 = tmp50 - tmp55
    tmp64 = tmp63 * tmp6
    tmp65 = tl_math.exp(tmp64)
    tmp66 = tmp62 + tmp65
    tmp67 = tmp54 - tmp55
    tmp68 = tmp67 * tmp6
    tmp69 = tl_math.exp(tmp68)
    tmp70 = tmp66 + tmp69
    tmp72 = tl.where(tmp2, tmp4, tmp71)
    tmp73 = tmp72 * tmp6
    tmp75 = tl.where(tmp9, tmp4, tmp74)
    tmp76 = tmp75 * tmp6
    tmp77 = triton_helpers.maximum(tmp73, tmp76)
    tmp79 = tl.where(tmp15, tmp4, tmp78)
    tmp80 = tmp79 * tmp6
    tmp81 = triton_helpers.maximum(tmp77, tmp80)
    tmp83 = tl.where(tmp21, tmp4, tmp82)
    tmp84 = tmp83 * tmp6
    tmp85 = triton_helpers.maximum(tmp81, tmp84)
    tmp86 = tmp73 - tmp85
    tmp87 = tmp86 * tmp6
    tmp88 = tl_math.exp(tmp87)
    tmp89 = tmp76 - tmp85
    tmp90 = tmp89 * tmp6
    tmp91 = tl_math.exp(tmp90)
    tmp92 = tmp88 + tmp91
    tmp93 = tmp80 - tmp85
    tmp94 = tmp93 * tmp6
    tmp95 = tl_math.exp(tmp94)
    tmp96 = tmp92 + tmp95
    tmp97 = tmp84 - tmp85
    tmp98 = tmp97 * tmp6
    tmp99 = tl_math.exp(tmp98)
    tmp100 = tmp96 + tmp99
    tmp102 = tl.where(tmp2, tmp4, tmp101)
    tmp103 = tmp102 * tmp6
    tmp105 = tl.where(tmp9, tmp4, tmp104)
    tmp106 = tmp105 * tmp6
    tmp107 = triton_helpers.maximum(tmp103, tmp106)
    tmp109 = tl.where(tmp15, tmp4, tmp108)
    tmp110 = tmp109 * tmp6
    tmp111 = triton_helpers.maximum(tmp107, tmp110)
    tmp113 = tl.where(tmp21, tmp4, tmp112)
    tmp114 = tmp113 * tmp6
    tmp115 = triton_helpers.maximum(tmp111, tmp114)
    tmp116 = tmp103 - tmp115
    tmp117 = tmp116 * tmp6
    tmp118 = tl_math.exp(tmp117)
    tmp119 = tmp106 - tmp115
    tmp120 = tmp119 * tmp6
    tmp121 = tl_math.exp(tmp120)
    tmp122 = tmp118 + tmp121
    tmp123 = tmp110 - tmp115
    tmp124 = tmp123 * tmp6
    tmp125 = tl_math.exp(tmp124)
    tmp126 = tmp122 + tmp125
    tmp127 = tmp114 - tmp115
    tmp128 = tmp127 * tmp6
    tmp129 = tl_math.exp(tmp128)
    tmp130 = tmp126 + tmp129
    tl.store(out_ptr0 + x2, tmp25, xmask)
    tl.store(out_ptr1 + x2, tmp40, xmask)
    tl.store(out_ptr2 + x2, tmp55, xmask)
    tl.store(out_ptr3 + x2, tmp70, xmask)
    tl.store(out_ptr4 + x2, tmp85, xmask)
    tl.store(out_ptr5 + x2, tmp100, xmask)
    tl.store(out_ptr6 + x2, tmp115, xmask)
    tl.store(out_ptr7 + x2, tmp130, xmask)


@triton.jit
def triton_poi_fused__softmax_eq_masked_fill_11(in_out_ptr0, in_out_ptr1,
    in_out_ptr2, in_out_ptr3, in_ptr0, in_ptr1, in_ptr2, in_ptr3, in_ptr4,
    in_ptr5, in_ptr6, in_ptr7, in_ptr8, xnumel, XBLOCK: tl.constexpr):
    xnumel = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x2 = xindex // 16
    x3 = xindex
    x4 = xindex // 4
    tmp0 = tl.load(in_ptr0 + (x0 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp3 = tl.load(in_out_ptr0 + x3, xmask)
    tmp8 = tl.load(in_ptr1 + x4, xmask, eviction_policy='evict_last')
    tmp12 = tl.load(in_ptr2 + x4, xmask, eviction_policy='evict_last')
    tmp14 = tl.load(in_out_ptr1 + x3, xmask)
    tmp17 = tl.load(in_ptr3 + x4, xmask, eviction_policy='evict_last')
    tmp21 = tl.load(in_ptr4 + x4, xmask, eviction_policy='evict_last')
    tmp23 = tl.load(in_out_ptr2 + x3, xmask)
    tmp26 = tl.load(in_ptr5 + x4, xmask, eviction_policy='evict_last')
    tmp30 = tl.load(in_ptr6 + x4, xmask, eviction_policy='evict_last')
    tmp32 = tl.load(in_out_ptr3 + x3, xmask)
    tmp35 = tl.load(in_ptr7 + x4, xmask, eviction_policy='evict_last')
    tmp39 = tl.load(in_ptr8 + x4, xmask, eviction_policy='evict_last')
    tmp1 = 0.0
    tmp2 = tmp0 == tmp1
    tmp4 = -1000000000.0
    tmp5 = tl.where(tmp2, tmp4, tmp3)
    tmp6 = 1.0
    tmp7 = tmp5 * tmp6
    tmp9 = tmp7 - tmp8
    tmp10 = tmp9 * tmp6
    tmp11 = tl_math.exp(tmp10)
    tmp13 = tmp11 / tmp12
    tmp15 = tl.where(tmp2, tmp4, tmp14)
    tmp16 = tmp15 * tmp6
    tmp18 = tmp16 - tmp17
    tmp19 = tmp18 * tmp6
    tmp20 = tl_math.exp(tmp19)
    tmp22 = tmp20 / tmp21
    tmp24 = tl.where(tmp2, tmp4, tmp23)
    tmp25 = tmp24 * tmp6
    tmp27 = tmp25 - tmp26
    tmp28 = tmp27 * tmp6
    tmp29 = tl_math.exp(tmp28)
    tmp31 = tmp29 / tmp30
    tmp33 = tl.where(tmp2, tmp4, tmp32)
    tmp34 = tmp33 * tmp6
    tmp36 = tmp34 - tmp35
    tmp37 = tmp36 * tmp6
    tmp38 = tl_math.exp(tmp37)
    tmp40 = tmp38 / tmp39
    tl.store(in_out_ptr0 + x3, tmp13, xmask)
    tl.store(in_out_ptr1 + x3, tmp22, xmask)
    tl.store(in_out_ptr2 + x3, tmp31, xmask)
    tl.store(in_out_ptr3 + x3, tmp40, xmask)


@triton.jit
def triton_poi_fused__softmax_eq_masked_fill_12(in_ptr0, in_ptr1, in_ptr2,
    in_ptr3, out_ptr0, out_ptr1, out_ptr2, out_ptr3, out_ptr4, out_ptr5,
    xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4
    x2 = xindex
    tmp0 = tl.load(in_ptr0 + 4 * x1, xmask, eviction_policy='evict_last')
    tmp3 = tl.load(in_ptr1 + 4 * x2, xmask, eviction_policy='evict_last')
    tmp8 = tl.load(in_ptr0 + (1 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp10 = tl.load(in_ptr1 + (1 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp14 = tl.load(in_ptr0 + (2 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp16 = tl.load(in_ptr1 + (2 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp20 = tl.load(in_ptr0 + (3 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp22 = tl.load(in_ptr1 + (3 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp41 = tl.load(in_ptr2 + 4 * x2, xmask, eviction_policy='evict_last')
    tmp44 = tl.load(in_ptr2 + (1 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp48 = tl.load(in_ptr2 + (2 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp52 = tl.load(in_ptr2 + (3 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp71 = tl.load(in_ptr3 + 4 * x2, xmask, eviction_policy='evict_last')
    tmp74 = tl.load(in_ptr3 + (1 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp78 = tl.load(in_ptr3 + (2 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp82 = tl.load(in_ptr3 + (3 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp1 = 0.0
    tmp2 = tmp0 == tmp1
    tmp4 = -1000000000.0
    tmp5 = tl.where(tmp2, tmp4, tmp3)
    tmp6 = 1.0
    tmp7 = tmp5 * tmp6
    tmp9 = tmp8 == tmp1
    tmp11 = tl.where(tmp9, tmp4, tmp10)
    tmp12 = tmp11 * tmp6
    tmp13 = triton_helpers.maximum(tmp7, tmp12)
    tmp15 = tmp14 == tmp1
    tmp17 = tl.where(tmp15, tmp4, tmp16)
    tmp18 = tmp17 * tmp6
    tmp19 = triton_helpers.maximum(tmp13, tmp18)
    tmp21 = tmp20 == tmp1
    tmp23 = tl.where(tmp21, tmp4, tmp22)
    tmp24 = tmp23 * tmp6
    tmp25 = triton_helpers.maximum(tmp19, tmp24)
    tmp26 = tmp7 - tmp25
    tmp27 = tmp26 * tmp6
    tmp28 = tl_math.exp(tmp27)
    tmp29 = tmp12 - tmp25
    tmp30 = tmp29 * tmp6
    tmp31 = tl_math.exp(tmp30)
    tmp32 = tmp28 + tmp31
    tmp33 = tmp18 - tmp25
    tmp34 = tmp33 * tmp6
    tmp35 = tl_math.exp(tmp34)
    tmp36 = tmp32 + tmp35
    tmp37 = tmp24 - tmp25
    tmp38 = tmp37 * tmp6
    tmp39 = tl_math.exp(tmp38)
    tmp40 = tmp36 + tmp39
    tmp42 = tl.where(tmp2, tmp4, tmp41)
    tmp43 = tmp42 * tmp6
    tmp45 = tl.where(tmp9, tmp4, tmp44)
    tmp46 = tmp45 * tmp6
    tmp47 = triton_helpers.maximum(tmp43, tmp46)
    tmp49 = tl.where(tmp15, tmp4, tmp48)
    tmp50 = tmp49 * tmp6
    tmp51 = triton_helpers.maximum(tmp47, tmp50)
    tmp53 = tl.where(tmp21, tmp4, tmp52)
    tmp54 = tmp53 * tmp6
    tmp55 = triton_helpers.maximum(tmp51, tmp54)
    tmp56 = tmp43 - tmp55
    tmp57 = tmp56 * tmp6
    tmp58 = tl_math.exp(tmp57)
    tmp59 = tmp46 - tmp55
    tmp60 = tmp59 * tmp6
    tmp61 = tl_math.exp(tmp60)
    tmp62 = tmp58 + tmp61
    tmp63 = tmp50 - tmp55
    tmp64 = tmp63 * tmp6
    tmp65 = tl_math.exp(tmp64)
    tmp66 = tmp62 + tmp65
    tmp67 = tmp54 - tmp55
    tmp68 = tmp67 * tmp6
    tmp69 = tl_math.exp(tmp68)
    tmp70 = tmp66 + tmp69
    tmp72 = tl.where(tmp2, tmp4, tmp71)
    tmp73 = tmp72 * tmp6
    tmp75 = tl.where(tmp9, tmp4, tmp74)
    tmp76 = tmp75 * tmp6
    tmp77 = triton_helpers.maximum(tmp73, tmp76)
    tmp79 = tl.where(tmp15, tmp4, tmp78)
    tmp80 = tmp79 * tmp6
    tmp81 = triton_helpers.maximum(tmp77, tmp80)
    tmp83 = tl.where(tmp21, tmp4, tmp82)
    tmp84 = tmp83 * tmp6
    tmp85 = triton_helpers.maximum(tmp81, tmp84)
    tmp86 = tmp73 - tmp85
    tmp87 = tmp86 * tmp6
    tmp88 = tl_math.exp(tmp87)
    tmp89 = tmp76 - tmp85
    tmp90 = tmp89 * tmp6
    tmp91 = tl_math.exp(tmp90)
    tmp92 = tmp88 + tmp91
    tmp93 = tmp80 - tmp85
    tmp94 = tmp93 * tmp6
    tmp95 = tl_math.exp(tmp94)
    tmp96 = tmp92 + tmp95
    tmp97 = tmp84 - tmp85
    tmp98 = tmp97 * tmp6
    tmp99 = tl_math.exp(tmp98)
    tmp100 = tmp96 + tmp99
    tl.store(out_ptr0 + x2, tmp25, xmask)
    tl.store(out_ptr1 + x2, tmp40, xmask)
    tl.store(out_ptr2 + x2, tmp55, xmask)
    tl.store(out_ptr3 + x2, tmp70, xmask)
    tl.store(out_ptr4 + x2, tmp85, xmask)
    tl.store(out_ptr5 + x2, tmp100, xmask)


@triton.jit
def triton_poi_fused_bmm_13(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (8 + 12 * x0), xmask, eviction_policy='evict_last'
        )
    tl.store(out_ptr0 + x0, tmp0, xmask)


@triton.jit
def triton_poi_fused_bmm_14(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (9 + 12 * x0), xmask, eviction_policy='evict_last'
        )
    tl.store(out_ptr0 + x0, tmp0, xmask)


@triton.jit
def triton_poi_fused_bmm_15(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (10 + 12 * x0), xmask, eviction_policy=
        'evict_last')
    tl.store(out_ptr0 + x0, tmp0, xmask)


@triton.jit
def triton_poi_fused_bmm_16(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (11 + 12 * x0), xmask, eviction_policy=
        'evict_last')
    tl.store(out_ptr0 + x0, tmp0, xmask)


@triton.jit
def triton_poi_fused_cat_17(in_ptr0, in_ptr1, in_ptr2, in_ptr3, 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
    tmp0 = x0
    tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 3, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.full([1], 2, tl.int64)
    tmp6 = tmp0 < tmp5
    tmp7 = tmp6 & tmp4
    tmp8 = tl.full([1], 1, tl.int64)
    tmp9 = tmp0 < tmp8
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + x1, tmp10 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp12 = tmp0 >= tmp8
    tmp13 = tmp12 & tmp7
    tmp14 = tl.load(in_ptr1 + x1, tmp13 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp15 = tl.where(tmp9, tmp11, tmp14)
    tmp16 = tl.full(tmp15.shape, 0.0, tmp15.dtype)
    tmp17 = tl.where(tmp7, tmp15, tmp16)
    tmp18 = tmp0 >= tmp5
    tmp19 = tmp18 & tmp4
    tmp20 = tl.load(in_ptr2 + x1, tmp19 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp21 = tl.where(tmp6, tmp17, tmp20)
    tmp22 = tl.full(tmp21.shape, 0.0, tmp21.dtype)
    tmp23 = tl.where(tmp4, tmp21, tmp22)
    tmp24 = tmp0 >= tmp3
    tl.full([1], 4, tl.int64)
    tmp27 = tl.load(in_ptr3 + x1, tmp24 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp28 = tl.where(tmp4, tmp23, tmp27)
    tl.store(out_ptr0 + (x0 + 8 * x1), tmp28, xmask)


@triton.jit
def triton_poi_fused_cat_18(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
    x2 = xindex
    x0 = xindex % 4
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tl.store(out_ptr0 + (x0 + 8 * x1), tmp0, xmask)


@triton.jit
def triton_poi_fused__softmax_div_eq_masked_fill_mul_relu_sum_19(in_out_ptr0,
    in_ptr0, in_ptr1, in_ptr2, in_ptr3, out_ptr0, out_ptr1, xnumel, XBLOCK:
    tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4
    x2 = xindex
    x0 = xindex % 4
    tmp0 = tl.load(in_ptr0 + 4 * x1, xmask, eviction_policy='evict_last')
    tmp3 = tl.load(in_ptr1 + 4 * x2, xmask, eviction_policy='evict_last')
    tmp8 = tl.load(in_ptr0 + (1 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp10 = tl.load(in_ptr1 + (1 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp14 = tl.load(in_ptr0 + (2 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp16 = tl.load(in_ptr1 + (2 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp20 = tl.load(in_ptr0 + (3 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp22 = tl.load(in_ptr1 + (3 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp41 = tl.load(in_ptr2 + (x0 + 16 * x1), xmask)
    tmp42 = tl.load(in_ptr3 + x0, xmask, eviction_policy='evict_last')
    tmp45 = tl.load(in_ptr2 + (4 + x0 + 16 * x1), xmask)
    tmp49 = tl.load(in_ptr2 + (8 + x0 + 16 * x1), xmask)
    tmp53 = tl.load(in_ptr2 + (12 + x0 + 16 * x1), xmask)
    tmp1 = 0.0
    tmp2 = tmp0 == tmp1
    tmp4 = -1000000000.0
    tmp5 = tl.where(tmp2, tmp4, tmp3)
    tmp6 = 1.0
    tmp7 = tmp5 * tmp6
    tmp9 = tmp8 == tmp1
    tmp11 = tl.where(tmp9, tmp4, tmp10)
    tmp12 = tmp11 * tmp6
    tmp13 = triton_helpers.maximum(tmp7, tmp12)
    tmp15 = tmp14 == tmp1
    tmp17 = tl.where(tmp15, tmp4, tmp16)
    tmp18 = tmp17 * tmp6
    tmp19 = triton_helpers.maximum(tmp13, tmp18)
    tmp21 = tmp20 == tmp1
    tmp23 = tl.where(tmp21, tmp4, tmp22)
    tmp24 = tmp23 * tmp6
    tmp25 = triton_helpers.maximum(tmp19, tmp24)
    tmp26 = tmp7 - tmp25
    tmp27 = tmp26 * tmp6
    tmp28 = tl_math.exp(tmp27)
    tmp29 = tmp12 - tmp25
    tmp30 = tmp29 * tmp6
    tmp31 = tl_math.exp(tmp30)
    tmp32 = tmp28 + tmp31
    tmp33 = tmp18 - tmp25
    tmp34 = tmp33 * tmp6
    tmp35 = tl_math.exp(tmp34)
    tmp36 = tmp32 + tmp35
    tmp37 = tmp24 - tmp25
    tmp38 = tmp37 * tmp6
    tmp39 = tl_math.exp(tmp38)
    tmp40 = tmp36 + tmp39
    tmp43 = tmp41 + tmp42
    tmp44 = tmp43 * tmp0
    tmp46 = tmp45 + tmp42
    tmp47 = tmp46 * tmp8
    tmp48 = tmp44 + tmp47
    tmp50 = tmp49 + tmp42
    tmp51 = tmp50 * tmp14
    tmp52 = tmp48 + tmp51
    tmp54 = tmp53 + tmp42
    tmp55 = tmp54 * tmp20
    tmp56 = tmp52 + tmp55
    tmp57 = tmp0 + tmp8
    tmp58 = tmp57 + tmp14
    tmp59 = tmp58 + tmp20
    tmp60 = tmp56 / tmp59
    tmp61 = tl.full([1], 0, tl.int32)
    tmp62 = triton_helpers.maximum(tmp61, tmp60)
    tl.store(out_ptr0 + x2, tmp25, xmask)
    tl.store(out_ptr1 + x2, tmp40, xmask)
    tl.store(in_out_ptr0 + x2, tmp62, xmask)


@triton.jit
def triton_poi_fused_div_mul_relu_sum_20(in_out_ptr0, in_ptr0, in_ptr1,
    in_ptr2, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x1 = xindex // 4
    x2 = xindex
    tmp0 = tl.load(in_ptr0 + (x0 + 16 * x1), xmask)
    tmp1 = tl.load(in_ptr1 + x0, xmask, eviction_policy='evict_last')
    tmp3 = tl.load(in_ptr2 + 4 * x1, xmask, eviction_policy='evict_last')
    tmp5 = tl.load(in_ptr0 + (4 + x0 + 16 * x1), xmask)
    tmp7 = tl.load(in_ptr2 + (1 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp10 = tl.load(in_ptr0 + (8 + x0 + 16 * x1), xmask)
    tmp12 = tl.load(in_ptr2 + (2 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp15 = tl.load(in_ptr0 + (12 + x0 + 16 * x1), xmask)
    tmp17 = tl.load(in_ptr2 + (3 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp2 = tmp0 + tmp1
    tmp4 = tmp2 * tmp3
    tmp6 = tmp5 + tmp1
    tmp8 = tmp6 * tmp7
    tmp9 = tmp4 + tmp8
    tmp11 = tmp10 + tmp1
    tmp13 = tmp11 * tmp12
    tmp14 = tmp9 + tmp13
    tmp16 = tmp15 + tmp1
    tmp18 = tmp16 * tmp17
    tmp19 = tmp14 + tmp18
    tmp20 = tmp3 + tmp7
    tmp21 = tmp20 + tmp12
    tmp22 = tmp21 + tmp17
    tmp23 = tmp19 / tmp22
    tmp24 = tl.full([1], 0, tl.int32)
    tmp25 = triton_helpers.maximum(tmp24, tmp23)
    tl.store(in_out_ptr0 + x2, tmp25, xmask)


@triton.jit
def triton_poi_fused_relu_threshold_backward_21(in_ptr0, in_ptr1, out_ptr0,
    out_ptr1, xnumel, XBLOCK: tl.constexpr):
    xnumel = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x0 = xindex % 4
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp1 = tl.load(in_ptr1 + x0, xmask, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tmp3 = tl.full([1], 0, tl.int32)
    tmp4 = triton_helpers.maximum(tmp3, tmp2)
    tmp5 = 0.0
    tmp6 = tmp4 <= tmp5
    tl.store(out_ptr0 + x2, tmp4, xmask)
    tl.store(out_ptr1 + x2, tmp6, xmask)


@triton.jit
def triton_poi_fused_add_mul_22(in_ptr0, in_ptr1, out_ptr0, out_ptr1,
    xnumel, XBLOCK: tl.constexpr):
    xnumel = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x2 = xindex // 16
    x3 = xindex // 4
    x4 = xindex
    tmp0 = tl.load(in_ptr0 + (x0 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp4 = tl.load(in_ptr1 + (4 + x0 + 12 * x3), xmask)
    tmp6 = tl.load(in_ptr1 + (x0 + 12 * x3), xmask)
    tmp1 = tl.sigmoid(tmp0)
    tmp2 = 1.0
    tmp3 = tmp1 + tmp2
    tmp5 = tmp3 * tmp4
    tmp7 = tmp3 * tmp6
    tl.store(out_ptr0 + x4, tmp5, xmask)
    tl.store(out_ptr1 + x4, tmp7, xmask)


@triton.jit
def triton_poi_fused_bmm_23(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + 4 * x0, xmask, eviction_policy='evict_last')
    tl.store(out_ptr0 + x0, tmp0, xmask)


@triton.jit
def triton_poi_fused_bmm_24(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (1 + 4 * x0), xmask, eviction_policy='evict_last')
    tl.store(out_ptr0 + x0, tmp0, xmask)


@triton.jit
def triton_poi_fused_bmm_25(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (2 + 4 * x0), xmask, eviction_policy='evict_last')
    tl.store(out_ptr0 + x0, tmp0, xmask)


@triton.jit
def triton_poi_fused_bmm_26(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (3 + 4 * x0), xmask, eviction_policy='evict_last')
    tl.store(out_ptr0 + x0, tmp0, xmask)


@triton.jit
def triton_poi_fused_add_cat_27(in_out_ptr0, in_ptr0, in_ptr1, in_ptr2,
    in_ptr3, in_ptr4, in_ptr5, 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
    x2 = xindex
    tmp29 = tl.load(in_ptr4 + x2, xmask)
    tmp30 = tl.load(in_ptr5 + x0, xmask, eviction_policy='evict_last')
    tmp0 = x0
    tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 3, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.full([1], 2, tl.int64)
    tmp6 = tmp0 < tmp5
    tmp7 = tmp6 & tmp4
    tmp8 = tl.full([1], 1, tl.int64)
    tmp9 = tmp0 < tmp8
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + x1, tmp10 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp12 = tmp0 >= tmp8
    tmp13 = tmp12 & tmp7
    tmp14 = tl.load(in_ptr1 + x1, tmp13 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp15 = tl.where(tmp9, tmp11, tmp14)
    tmp16 = tl.full(tmp15.shape, 0.0, tmp15.dtype)
    tmp17 = tl.where(tmp7, tmp15, tmp16)
    tmp18 = tmp0 >= tmp5
    tmp19 = tmp18 & tmp4
    tmp20 = tl.load(in_ptr2 + x1, tmp19 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp21 = tl.where(tmp6, tmp17, tmp20)
    tmp22 = tl.full(tmp21.shape, 0.0, tmp21.dtype)
    tmp23 = tl.where(tmp4, tmp21, tmp22)
    tmp24 = tmp0 >= tmp3
    tl.full([1], 4, tl.int64)
    tmp27 = tl.load(in_ptr3 + x1, tmp24 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp28 = tl.where(tmp4, tmp23, tmp27)
    tmp31 = tmp29 + tmp30
    tmp32 = tmp31 + tmp28
    tl.store(in_out_ptr0 + x2, tmp32, xmask)


@triton.jit
def triton_poi_fused_add_cat_28(in_out_ptr0, in_ptr0, in_ptr1, in_ptr2,
    in_ptr3, in_ptr4, 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
    x2 = xindex
    tmp29 = tl.load(in_out_ptr0 + x2, xmask)
    tmp30 = tl.load(in_ptr4 + x0, xmask, eviction_policy='evict_last')
    tmp0 = x0
    tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 3, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.full([1], 2, tl.int64)
    tmp6 = tmp0 < tmp5
    tmp7 = tmp6 & tmp4
    tmp8 = tl.full([1], 1, tl.int64)
    tmp9 = tmp0 < tmp8
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + x1, tmp10 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp12 = tmp0 >= tmp8
    tmp13 = tmp12 & tmp7
    tmp14 = tl.load(in_ptr1 + x1, tmp13 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp15 = tl.where(tmp9, tmp11, tmp14)
    tmp16 = tl.full(tmp15.shape, 0.0, tmp15.dtype)
    tmp17 = tl.where(tmp7, tmp15, tmp16)
    tmp18 = tmp0 >= tmp5
    tmp19 = tmp18 & tmp4
    tmp20 = tl.load(in_ptr2 + x1, tmp19 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp21 = tl.where(tmp6, tmp17, tmp20)
    tmp22 = tl.full(tmp21.shape, 0.0, tmp21.dtype)
    tmp23 = tl.where(tmp4, tmp21, tmp22)
    tmp24 = tmp0 >= tmp3
    tl.full([1], 4, tl.int64)
    tmp27 = tl.load(in_ptr3 + x1, tmp24 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp28 = tl.where(tmp4, tmp23, tmp27)
    tmp31 = tmp29 + tmp30
    tmp32 = tmp31 + tmp28
    tl.store(in_out_ptr0 + x2, tmp32, xmask)


def call(args):
    (primals_1, primals_2, primals_3, primals_4, primals_5, primals_6,
        primals_7, primals_8, primals_9, primals_10, primals_11, primals_12,
        primals_13, primals_14, primals_15, primals_16, primals_17,
        primals_18, primals_19, primals_20, primals_21, primals_22,
        primals_23, primals_24, primals_25, primals_26, primals_27, primals_28
        ) = args
    args.clear()
    assert_size_stride(primals_1, (4, 4, 4), (16, 4, 1))
    assert_size_stride(primals_2, (4, 4), (4, 1))
    assert_size_stride(primals_3, (4,), (1,))
    assert_size_stride(primals_4, (4, 4, 4), (16, 4, 1))
    assert_size_stride(primals_5, (4, 4), (4, 1))
    assert_size_stride(primals_6, (4,), (1,))
    assert_size_stride(primals_7, (4, 4), (4, 1))
    assert_size_stride(primals_8, (4, 4), (4, 1))
    assert_size_stride(primals_9, (12, 4), (4, 1))
    assert_size_stride(primals_10, (12,), (1,))
    assert_size_stride(primals_11, (12, 4), (4, 1))
    assert_size_stride(primals_12, (12,), (1,))
    assert_size_stride(primals_13, (4, 8), (8, 1))
    assert_size_stride(primals_14, (4,), (1,))
    assert_size_stride(primals_15, (4, 8), (8, 1))
    assert_size_stride(primals_16, (4,), (1,))
    assert_size_stride(primals_17, (4, 4), (4, 1))
    assert_size_stride(primals_18, (4,), (1,))
    assert_size_stride(primals_19, (4, 4), (4, 1))
    assert_size_stride(primals_20, (4,), (1,))
    assert_size_stride(primals_21, (12, 4), (4, 1))
    assert_size_stride(primals_22, (12,), (1,))
    assert_size_stride(primals_23, (12, 4), (4, 1))
    assert_size_stride(primals_24, (12,), (1,))
    assert_size_stride(primals_25, (4, 4), (4, 1))
    assert_size_stride(primals_26, (4,), (1,))
    assert_size_stride(primals_27, (4, 4), (4, 1))
    assert_size_stride(primals_28, (4,), (1,))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((16, 4), (4, 1), torch.float32)
        extern_kernels.addmm(primals_3, reinterpret_tensor(primals_1, (16, 
            4), (4, 1), 0), reinterpret_tensor(primals_2, (4, 4), (1, 4), 0
            ), alpha=1, beta=1, out=buf0)
        del primals_2
        del primals_3
        buf1 = empty_strided_cuda((16, 4), (4, 1), torch.float32)
        extern_kernels.addmm(primals_6, reinterpret_tensor(primals_4, (16, 
            4), (4, 1), 0), reinterpret_tensor(primals_5, (4, 4), (1, 4), 0
            ), alpha=1, beta=1, out=buf1)
        del primals_5
        del primals_6
        buf2 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf673 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.bool)
        get_raw_stream(0)
        triton_poi_fused_relu_threshold_backward_0[grid(64)](buf0, buf2,
            buf673, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf3 = empty_strided_cuda((16, 12), (12, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(buf2, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_9, (4, 12), (1, 4), 0), out=buf3)
        buf4 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf672 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.bool)
        triton_poi_fused_relu_threshold_backward_0[grid(64)](buf1, buf4,
            buf672, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf5 = empty_strided_cuda((16, 12), (12, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(buf4, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_11, (4, 12), (1, 4), 0), out=buf5)
        buf6 = reinterpret_tensor(buf3, (4, 4, 12), (48, 12, 1), 0)
        del buf3
        triton_poi_fused_mul_1[grid(192)](buf6, primals_10, primals_7, 192,
            XBLOCK=256, num_warps=4, num_stages=1)
        buf7 = reinterpret_tensor(buf5, (4, 4, 12), (48, 12, 1), 0)
        del buf5
        triton_poi_fused_mul_1[grid(192)](buf7, primals_12, primals_8, 192,
            XBLOCK=256, num_warps=4, num_stages=1)
        buf8 = empty_strided_cuda((4, 4, 1), (4, 1, 16), torch.float32)
        triton_poi_fused_bmm_2[grid(16)](buf6, buf8, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf9 = empty_strided_cuda((4, 1, 4), (4, 16, 1), torch.float32)
        triton_poi_fused_bmm_3[grid(16)](buf7, buf9, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf10 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf8, buf9, out=buf10)
        buf11 = reinterpret_tensor(buf9, (4, 4, 1), (4, 1, 16), 0)
        del buf9
        triton_poi_fused_bmm_2[grid(16)](buf7, buf11, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf12 = reinterpret_tensor(buf8, (4, 1, 4), (4, 16, 1), 0)
        del buf8
        triton_poi_fused_bmm_3[grid(16)](buf6, buf12, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf13 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf11, buf12, out=buf13)
        buf24 = reinterpret_tensor(buf12, (4, 4, 1), (4, 1, 16), 0)
        del buf12
        triton_poi_fused_bmm_4[grid(16)](buf6, buf24, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf25 = reinterpret_tensor(buf11, (4, 1, 4), (4, 16, 1), 0)
        del buf11
        triton_poi_fused_bmm_5[grid(16)](buf7, buf25, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf26 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf24, buf25, out=buf26)
        buf40 = reinterpret_tensor(buf25, (4, 4, 1), (4, 1, 16), 0)
        del buf25
        triton_poi_fused_bmm_6[grid(16)](buf6, buf40, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf41 = reinterpret_tensor(buf24, (4, 1, 4), (4, 16, 1), 0)
        del buf24
        triton_poi_fused_bmm_7[grid(16)](buf7, buf41, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf42 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf40, buf41, out=buf42)
        buf56 = reinterpret_tensor(buf41, (4, 4, 1), (4, 1, 16), 0)
        del buf41
        triton_poi_fused_bmm_8[grid(16)](buf6, buf56, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf57 = reinterpret_tensor(buf40, (4, 1, 4), (4, 16, 1), 0)
        del buf40
        triton_poi_fused_bmm_9[grid(16)](buf7, buf57, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf58 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf56, buf57, out=buf58)
        buf14 = reinterpret_tensor(buf57, (4, 4, 1), (4, 1, 16), 0)
        del buf57
        buf15 = buf56
        del buf56
        buf30 = empty_strided_cuda((4, 4, 1), (4, 1, 16), torch.float32)
        buf31 = empty_strided_cuda((4, 4, 1), (4, 1, 16), torch.float32)
        buf46 = empty_strided_cuda((4, 4, 1), (4, 1, 16), torch.float32)
        buf47 = empty_strided_cuda((4, 4, 1), (4, 1, 16), torch.float32)
        buf62 = empty_strided_cuda((4, 4, 1), (4, 1, 16), torch.float32)
        buf63 = empty_strided_cuda((4, 4, 1), (4, 1, 16), torch.float32)
        triton_poi_fused__softmax_eq_masked_fill_10[grid(16)](primals_8,
            buf10, buf26, buf42, buf58, buf14, buf15, buf30, buf31, buf46,
            buf47, buf62, buf63, 16, XBLOCK=16, num_warps=1, num_stages=1)
        buf16 = buf10
        del buf10
        buf32 = buf26
        del buf26
        buf48 = buf42
        del buf42
        buf64 = buf58
        del buf58
        triton_poi_fused__softmax_eq_masked_fill_11[grid(64)](buf16, buf32,
            buf48, buf64, primals_8, buf14, buf15, buf30, buf31, buf46,
            buf47, buf62, buf63, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf27 = buf63
        del buf63
        triton_poi_fused_bmm_4[grid(16)](buf7, buf27, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf28 = reinterpret_tensor(buf62, (4, 1, 4), (4, 16, 1), 0)
        del buf62
        triton_poi_fused_bmm_5[grid(16)](buf6, buf28, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf29 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf27, buf28, out=buf29)
        buf43 = reinterpret_tensor(buf28, (4, 4, 1), (4, 1, 16), 0)
        del buf28
        triton_poi_fused_bmm_6[grid(16)](buf7, buf43, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf44 = reinterpret_tensor(buf27, (4, 1, 4), (4, 16, 1), 0)
        del buf27
        triton_poi_fused_bmm_7[grid(16)](buf6, buf44, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf45 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf43, buf44, out=buf45)
        buf17 = reinterpret_tensor(buf44, (4, 4, 1), (4, 1, 16), 0)
        del buf44
        buf18 = buf43
        del buf43
        buf33 = buf47
        del buf47
        buf34 = buf46
        del buf46
        buf49 = buf31
        del buf31
        buf50 = buf30
        del buf30
        triton_poi_fused__softmax_eq_masked_fill_12[grid(16)](primals_7,
            buf13, buf29, buf45, buf17, buf18, buf33, buf34, buf49, buf50, 
            16, XBLOCK=16, num_warps=1, num_stages=1)
        buf59 = buf15
        del buf15
        triton_poi_fused_bmm_8[grid(16)](buf7, buf59, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf60 = reinterpret_tensor(buf14, (4, 1, 4), (4, 16, 1), 0)
        del buf14
        triton_poi_fused_bmm_9[grid(16)](buf6, buf60, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf61 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf59, buf60, out=buf61)
        buf20 = reinterpret_tensor(buf60, (4, 4, 1), (4, 1, 16), 0)
        del buf60
        triton_poi_fused_bmm_13[grid(16)](buf7, buf20, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf21 = reinterpret_tensor(buf59, (4, 4, 1), (4, 1, 1), 0)
        del buf59
        extern_kernels.bmm(buf16, buf20, out=buf21)
        buf36 = buf20
        del buf20
        triton_poi_fused_bmm_14[grid(16)](buf7, buf36, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf37 = empty_strided_cuda((4, 4, 1), (4, 1, 1), torch.float32)
        extern_kernels.bmm(buf32, buf36, out=buf37)
        buf52 = buf36
        del buf36
        triton_poi_fused_bmm_15[grid(16)](buf7, buf52, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf53 = empty_strided_cuda((4, 4, 1), (4, 1, 1), torch.float32)
        extern_kernels.bmm(buf48, buf52, out=buf53)
        buf68 = buf52
        del buf52
        triton_poi_fused_bmm_16[grid(16)](buf7, buf68, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf69 = empty_strided_cuda((4, 4, 1), (4, 1, 1), torch.float32)
        extern_kernels.bmm(buf64, buf68, out=buf69)
        buf75 = empty_strided_cuda((4, 4, 8), (32, 8, 1), torch.float32)
        buf70 = reinterpret_tensor(buf75, (4, 4, 4), (32, 8, 1), 4)
        triton_poi_fused_cat_17[grid(64)](buf21, buf37, buf53, buf69, buf70,
            64, XBLOCK=64, num_warps=1, num_stages=1)
        buf74 = reinterpret_tensor(buf75, (4, 4, 4), (32, 8, 1), 0)
        triton_poi_fused_cat_18[grid(64)](buf0, buf74, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf78 = buf0
        del buf0
        extern_kernels.mm(reinterpret_tensor(buf75, (16, 8), (8, 1), 0),
            reinterpret_tensor(primals_13, (8, 4), (1, 8), 0), out=buf78)
        buf65 = reinterpret_tensor(buf69, (4, 4, 1), (4, 1, 16), 0)
        del buf69
        buf66 = reinterpret_tensor(buf53, (4, 4, 1), (4, 1, 16), 0)
        del buf53
        buf80 = reinterpret_tensor(buf37, (4, 4), (4, 1), 0)
        del buf37
        buf82 = buf80
        del buf80
        triton_poi_fused__softmax_div_eq_masked_fill_mul_relu_sum_19[grid(16)](
            buf82, primals_7, buf61, buf78, primals_14, buf65, buf66, 16,
            XBLOCK=16, num_warps=1, num_stages=1)
        buf19 = buf13
        del buf13
        buf35 = buf29
        del buf29
        buf51 = buf45
        del buf45
        buf67 = buf61
        del buf61
        triton_poi_fused__softmax_eq_masked_fill_11[grid(64)](buf19, buf35,
            buf51, buf67, primals_7, buf17, buf18, buf33, buf34, buf49,
            buf50, buf65, buf66, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf22 = buf66
        del buf66
        triton_poi_fused_bmm_13[grid(16)](buf6, buf22, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf23 = reinterpret_tensor(buf65, (4, 4, 1), (4, 1, 1), 0)
        del buf65
        extern_kernels.bmm(buf19, buf22, out=buf23)
        buf38 = buf22
        del buf22
        triton_poi_fused_bmm_14[grid(16)](buf6, buf38, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf39 = reinterpret_tensor(buf50, (4, 4, 1), (4, 1, 1), 0)
        del buf50
        extern_kernels.bmm(buf35, buf38, out=buf39)
        buf54 = buf38
        del buf38
        triton_poi_fused_bmm_15[grid(16)](buf6, buf54, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf55 = reinterpret_tensor(buf49, (4, 4, 1), (4, 1, 1), 0)
        del buf49
        extern_kernels.bmm(buf51, buf54, out=buf55)
        buf71 = buf54
        del buf54
        triton_poi_fused_bmm_16[grid(16)](buf6, buf71, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf72 = reinterpret_tensor(buf34, (4, 4, 1), (4, 1, 1), 0)
        del buf34
        extern_kernels.bmm(buf67, buf71, out=buf72)
        buf77 = empty_strided_cuda((4, 4, 8), (32, 8, 1), torch.float32)
        buf73 = reinterpret_tensor(buf77, (4, 4, 4), (32, 8, 1), 4)
        triton_poi_fused_cat_17[grid(64)](buf23, buf39, buf55, buf72, buf73,
            64, XBLOCK=64, num_warps=1, num_stages=1)
        buf76 = reinterpret_tensor(buf77, (4, 4, 4), (32, 8, 1), 0)
        triton_poi_fused_cat_18[grid(64)](buf1, buf76, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf79 = buf1
        del buf1
        extern_kernels.mm(reinterpret_tensor(buf77, (16, 8), (8, 1), 0),
            reinterpret_tensor(primals_15, (8, 4), (1, 8), 0), out=buf79)
        buf81 = reinterpret_tensor(buf72, (4, 4), (4, 1), 0)
        del buf72
        buf84 = buf81
        del buf81
        triton_poi_fused_div_mul_relu_sum_20[grid(16)](buf84, buf79,
            primals_16, primals_8, 16, XBLOCK=16, num_warps=1, num_stages=1)
        buf83 = reinterpret_tensor(buf55, (4, 4), (4, 1), 0)
        del buf55
        extern_kernels.addmm(primals_18, buf82, reinterpret_tensor(
            primals_17, (4, 4), (1, 4), 0), alpha=1, beta=1, out=buf83)
        buf85 = reinterpret_tensor(buf39, (4, 4), (4, 1), 0)
        del buf39
        extern_kernels.addmm(primals_20, buf84, reinterpret_tensor(
            primals_19, (4, 4), (1, 4), 0), alpha=1, beta=1, out=buf85)
        buf86 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf671 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.bool)
        triton_poi_fused_relu_threshold_backward_21[grid(64)](buf78,
            primals_14, buf86, buf671, 64, XBLOCK=64, num_warps=1, num_stages=1
            )
        buf87 = empty_strided_cuda((16, 12), (12, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(buf86, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_21, (4, 12), (1, 4), 0), out=buf87)
        buf88 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf670 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.bool)
        triton_poi_fused_relu_threshold_backward_21[grid(64)](buf79,
            primals_16, buf88, buf670, 64, XBLOCK=64, num_warps=1, num_stages=1
            )
        buf89 = empty_strided_cuda((16, 12), (12, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(buf88, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_23, (4, 12), (1, 4), 0), out=buf89)
        buf90 = reinterpret_tensor(buf87, (4, 4, 12), (48, 12, 1), 0)
        del buf87
        triton_poi_fused_mul_1[grid(192)](buf90, primals_22, primals_7, 192,
            XBLOCK=256, num_warps=4, num_stages=1)
        buf91 = reinterpret_tensor(buf89, (4, 4, 12), (48, 12, 1), 0)
        del buf89
        triton_poi_fused_mul_1[grid(192)](buf91, primals_24, primals_8, 192,
            XBLOCK=256, num_warps=4, num_stages=1)
        buf92 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf93 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        triton_poi_fused_add_mul_22[grid(64)](buf85, buf90, buf92, buf93, 
            64, XBLOCK=64, num_warps=1, num_stages=1)
        buf94 = reinterpret_tensor(buf23, (4, 4, 1), (4, 1, 16), 0)
        del buf23
        triton_poi_fused_bmm_23[grid(16)](buf92, buf94, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf95 = reinterpret_tensor(buf71, (4, 1, 4), (4, 16, 1), 0)
        del buf71
        triton_poi_fused_bmm_23[grid(16)](buf93, buf95, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf96 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf94, buf95, out=buf96)
        buf97 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf98 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        triton_poi_fused_add_mul_22[grid(64)](buf83, buf91, buf97, buf98, 
            64, XBLOCK=64, num_warps=1, num_stages=1)
        buf99 = reinterpret_tensor(buf95, (4, 4, 1), (4, 1, 16), 0)
        del buf95
        triton_poi_fused_bmm_23[grid(16)](buf97, buf99, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf100 = reinterpret_tensor(buf94, (4, 1, 4), (4, 16, 1), 0)
        del buf94
        triton_poi_fused_bmm_23[grid(16)](buf98, buf100, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf101 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf99, buf100, out=buf101)
        buf112 = buf99
        del buf99
        triton_poi_fused_bmm_24[grid(16)](buf92, buf112, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf113 = buf100
        del buf100
        triton_poi_fused_bmm_24[grid(16)](buf93, buf113, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf114 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf112, buf113, out=buf114)
        buf128 = reinterpret_tensor(buf113, (4, 4, 1), (4, 1, 16), 0)
        del buf113
        triton_poi_fused_bmm_25[grid(16)](buf92, buf128, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf129 = reinterpret_tensor(buf112, (4, 1, 4), (4, 16, 1), 0)
        del buf112
        triton_poi_fused_bmm_25[grid(16)](buf93, buf129, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf130 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf128, buf129, out=buf130)
        buf144 = reinterpret_tensor(buf129, (4, 4, 1), (4, 1, 16), 0)
        del buf129
        triton_poi_fused_bmm_26[grid(16)](buf92, buf144, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf145 = reinterpret_tensor(buf128, (4, 1, 4), (4, 16, 1), 0)
        del buf128
        triton_poi_fused_bmm_26[grid(16)](buf93, buf145, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf146 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf144, buf145, out=buf146)
        buf102 = reinterpret_tensor(buf145, (4, 4, 1), (4, 1, 16), 0)
        del buf145
        buf103 = buf144
        del buf144
        buf118 = buf33
        del buf33
        buf119 = buf18
        del buf18
        buf134 = buf17
        del buf17
        buf135 = reinterpret_tensor(buf21, (4, 4, 1), (4, 1, 16), 0)
        del buf21
        buf150 = buf68
        del buf68
        buf151 = empty_strided_cuda((4, 4, 1), (4, 1, 16), torch.float32)
        triton_poi_fused__softmax_eq_masked_fill_10[grid(16)](primals_7,
            buf96, buf114, buf130, buf146, buf102, buf103, buf118, buf119,
            buf134, buf135, buf150, buf151, 16, XBLOCK=16, num_warps=1,
            num_stages=1)
        buf104 = buf96
        del buf96
        buf120 = buf114
        del buf114
        buf136 = buf130
        del buf130
        buf152 = buf146
        del buf146
        triton_poi_fused__softmax_eq_masked_fill_11[grid(64)](buf104,
            buf120, buf136, buf152, primals_7, buf102, buf103, buf118,
            buf119, buf134, buf135, buf150, buf151, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf115 = buf151
        del buf151
        triton_poi_fused_bmm_24[grid(16)](buf97, buf115, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf116 = reinterpret_tensor(buf150, (4, 1, 4), (4, 16, 1), 0)
        del buf150
        triton_poi_fused_bmm_24[grid(16)](buf98, buf116, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf117 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf115, buf116, out=buf117)
        buf131 = reinterpret_tensor(buf116, (4, 4, 1), (4, 1, 16), 0)
        del buf116
        triton_poi_fused_bmm_25[grid(16)](buf97, buf131, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf132 = reinterpret_tensor(buf115, (4, 1, 4), (4, 16, 1), 0)
        del buf115
        triton_poi_fused_bmm_25[grid(16)](buf98, buf132, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf133 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf131, buf132, out=buf133)
        buf147 = reinterpret_tensor(buf132, (4, 4, 1), (4, 1, 16), 0)
        del buf132
        triton_poi_fused_bmm_26[grid(16)](buf97, buf147, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf148 = reinterpret_tensor(buf131, (4, 1, 4), (4, 16, 1), 0)
        del buf131
        triton_poi_fused_bmm_26[grid(16)](buf98, buf148, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf149 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf147, buf148, out=buf149)
        buf105 = reinterpret_tensor(buf148, (4, 4, 1), (4, 1, 16), 0)
        del buf148
        buf106 = buf147
        del buf147
        buf121 = buf135
        del buf135
        buf122 = buf134
        del buf134
        buf137 = buf119
        del buf119
        buf138 = buf118
        del buf118
        buf153 = buf103
        del buf103
        buf154 = buf102
        del buf102
        triton_poi_fused__softmax_eq_masked_fill_10[grid(16)](primals_8,
            buf101, buf117, buf133, buf149, buf105, buf106, buf121, buf122,
            buf137, buf138, buf153, buf154, 16, XBLOCK=16, num_warps=1,
            num_stages=1)
        buf107 = buf101
        del buf101
        buf123 = buf117
        del buf117
        buf139 = buf133
        del buf133
        buf155 = buf149
        del buf149
        triton_poi_fused__softmax_eq_masked_fill_11[grid(64)](buf107,
            buf123, buf139, buf155, primals_8, buf105, buf106, buf121,
            buf122, buf137, buf138, buf153, buf154, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf108 = buf154
        del buf154
        triton_poi_fused_bmm_13[grid(16)](buf90, buf108, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf109 = reinterpret_tensor(buf153, (4, 4, 1), (4, 1, 1), 0)
        del buf153
        extern_kernels.bmm(buf104, buf108, out=buf109)
        buf110 = buf108
        del buf108
        triton_poi_fused_bmm_13[grid(16)](buf91, buf110, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf111 = reinterpret_tensor(buf138, (4, 4, 1), (4, 1, 1), 0)
        del buf138
        extern_kernels.bmm(buf107, buf110, out=buf111)
        buf124 = buf110
        del buf110
        triton_poi_fused_bmm_14[grid(16)](buf90, buf124, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf125 = reinterpret_tensor(buf137, (4, 4, 1), (4, 1, 1), 0)
        del buf137
        extern_kernels.bmm(buf120, buf124, out=buf125)
        buf126 = buf124
        del buf124
        triton_poi_fused_bmm_14[grid(16)](buf91, buf126, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf127 = reinterpret_tensor(buf122, (4, 4, 1), (4, 1, 1), 0)
        del buf122
        extern_kernels.bmm(buf123, buf126, out=buf127)
        buf140 = buf126
        del buf126
        triton_poi_fused_bmm_15[grid(16)](buf90, buf140, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf141 = reinterpret_tensor(buf121, (4, 4, 1), (4, 1, 1), 0)
        del buf121
        extern_kernels.bmm(buf136, buf140, out=buf141)
        buf142 = buf140
        del buf140
        triton_poi_fused_bmm_15[grid(16)](buf91, buf142, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf143 = reinterpret_tensor(buf106, (4, 4, 1), (4, 1, 1), 0)
        del buf106
        extern_kernels.bmm(buf139, buf142, out=buf143)
        buf156 = buf142
        del buf142
        triton_poi_fused_bmm_16[grid(16)](buf90, buf156, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf157 = reinterpret_tensor(buf105, (4, 4, 1), (4, 1, 1), 0)
        del buf105
        extern_kernels.bmm(buf152, buf156, out=buf157)
        buf158 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf162 = buf158
        del buf158
        triton_poi_fused_add_cat_27[grid(64)](buf162, buf109, buf125,
            buf141, buf157, buf78, primals_14, 64, XBLOCK=64, num_warps=1,
            num_stages=1)
        buf159 = reinterpret_tensor(buf157, (4, 4, 1), (4, 1, 16), 0)
        del buf157
        triton_poi_fused_bmm_16[grid(16)](buf91, buf159, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf160 = buf141
        del buf141
        extern_kernels.bmm(buf155, buf159, out=buf160)
        buf161 = reinterpret_tensor(buf78, (4, 4, 4), (16, 4, 1), 0)
        del buf78
        buf164 = buf161
        del buf161
        triton_poi_fused_add_cat_27[grid(64)](buf164, buf111, buf127,
            buf143, buf160, buf79, primals_16, 64, XBLOCK=64, num_warps=1,
            num_stages=1)
        buf163 = buf79
        del buf79
        extern_kernels.addmm(primals_26, reinterpret_tensor(buf162, (16, 4),
            (4, 1), 0), reinterpret_tensor(primals_25, (4, 4), (1, 4), 0),
            alpha=1, beta=1, out=buf163)
        buf165 = empty_strided_cuda((16, 4), (4, 1), torch.float32)
        extern_kernels.addmm(primals_28, reinterpret_tensor(buf164, (16, 4),
            (4, 1), 0), reinterpret_tensor(primals_27, (4, 4), (1, 4), 0),
            alpha=1, beta=1, out=buf165)
        buf166 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf669 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.bool)
        triton_poi_fused_relu_threshold_backward_0[grid(64)](buf163, buf166,
            buf669, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf167 = empty_strided_cuda((16, 12), (12, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(buf166, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_9, (4, 12), (1, 4), 0), out=buf167)
        buf168 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf668 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.bool)
        triton_poi_fused_relu_threshold_backward_0[grid(64)](buf165, buf168,
            buf668, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf169 = empty_strided_cuda((16, 12), (12, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(buf168, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_11, (4, 12), (1, 4), 0), out=buf169)
        buf170 = reinterpret_tensor(buf167, (4, 4, 12), (48, 12, 1), 0)
        del buf167
        triton_poi_fused_mul_1[grid(192)](buf170, primals_10, primals_7, 
            192, XBLOCK=256, num_warps=4, num_stages=1)
        buf171 = reinterpret_tensor(buf169, (4, 4, 12), (48, 12, 1), 0)
        del buf169
        triton_poi_fused_mul_1[grid(192)](buf171, primals_12, primals_8, 
            192, XBLOCK=256, num_warps=4, num_stages=1)
        buf172 = reinterpret_tensor(buf160, (4, 4, 1), (4, 1, 16), 0)
        del buf160
        triton_poi_fused_bmm_2[grid(16)](buf170, buf172, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf173 = reinterpret_tensor(buf143, (4, 1, 4), (4, 16, 1), 0)
        del buf143
        triton_poi_fused_bmm_3[grid(16)](buf171, buf173, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf174 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf172, buf173, out=buf174)
        buf175 = reinterpret_tensor(buf173, (4, 4, 1), (4, 1, 16), 0)
        del buf173
        triton_poi_fused_bmm_2[grid(16)](buf171, buf175, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf176 = reinterpret_tensor(buf172, (4, 1, 4), (4, 16, 1), 0)
        del buf172
        triton_poi_fused_bmm_3[grid(16)](buf170, buf176, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf177 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf175, buf176, out=buf177)
        buf188 = reinterpret_tensor(buf176, (4, 4, 1), (4, 1, 16), 0)
        del buf176
        triton_poi_fused_bmm_4[grid(16)](buf170, buf188, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf189 = reinterpret_tensor(buf175, (4, 1, 4), (4, 16, 1), 0)
        del buf175
        triton_poi_fused_bmm_5[grid(16)](buf171, buf189, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf190 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf188, buf189, out=buf190)
        buf204 = reinterpret_tensor(buf189, (4, 4, 1), (4, 1, 16), 0)
        del buf189
        triton_poi_fused_bmm_6[grid(16)](buf170, buf204, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf205 = reinterpret_tensor(buf188, (4, 1, 4), (4, 16, 1), 0)
        del buf188
        triton_poi_fused_bmm_7[grid(16)](buf171, buf205, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf206 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf204, buf205, out=buf206)
        buf220 = reinterpret_tensor(buf205, (4, 4, 1), (4, 1, 16), 0)
        del buf205
        triton_poi_fused_bmm_8[grid(16)](buf170, buf220, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf221 = reinterpret_tensor(buf204, (4, 1, 4), (4, 16, 1), 0)
        del buf204
        triton_poi_fused_bmm_9[grid(16)](buf171, buf221, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf222 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf220, buf221, out=buf222)
        buf178 = reinterpret_tensor(buf221, (4, 4, 1), (4, 1, 16), 0)
        del buf221
        buf179 = buf220
        del buf220
        buf194 = reinterpret_tensor(buf127, (4, 4, 1), (4, 1, 16), 0)
        del buf127
        buf195 = reinterpret_tensor(buf111, (4, 4, 1), (4, 1, 16), 0)
        del buf111
        buf210 = buf159
        del buf159
        buf211 = reinterpret_tensor(buf125, (4, 4, 1), (4, 1, 16), 0)
        del buf125
        buf226 = reinterpret_tensor(buf109, (4, 4, 1), (4, 1, 16), 0)
        del buf109
        buf227 = buf156
        del buf156
        triton_poi_fused__softmax_eq_masked_fill_10[grid(16)](primals_8,
            buf174, buf190, buf206, buf222, buf178, buf179, buf194, buf195,
            buf210, buf211, buf226, buf227, 16, XBLOCK=16, num_warps=1,
            num_stages=1)
        buf180 = buf174
        del buf174
        buf196 = buf190
        del buf190
        buf212 = buf206
        del buf206
        buf228 = buf222
        del buf222
        triton_poi_fused__softmax_eq_masked_fill_11[grid(64)](buf180,
            buf196, buf212, buf228, primals_8, buf178, buf179, buf194,
            buf195, buf210, buf211, buf226, buf227, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf191 = buf227
        del buf227
        triton_poi_fused_bmm_4[grid(16)](buf171, buf191, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf192 = reinterpret_tensor(buf226, (4, 1, 4), (4, 16, 1), 0)
        del buf226
        triton_poi_fused_bmm_5[grid(16)](buf170, buf192, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf193 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf191, buf192, out=buf193)
        buf207 = reinterpret_tensor(buf192, (4, 4, 1), (4, 1, 16), 0)
        del buf192
        triton_poi_fused_bmm_6[grid(16)](buf171, buf207, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf208 = reinterpret_tensor(buf191, (4, 1, 4), (4, 16, 1), 0)
        del buf191
        triton_poi_fused_bmm_7[grid(16)](buf170, buf208, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf209 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf207, buf208, out=buf209)
        buf181 = reinterpret_tensor(buf208, (4, 4, 1), (4, 1, 16), 0)
        del buf208
        buf182 = buf207
        del buf207
        buf197 = buf211
        del buf211
        buf198 = buf210
        del buf210
        buf213 = buf195
        del buf195
        buf214 = buf194
        del buf194
        triton_poi_fused__softmax_eq_masked_fill_12[grid(16)](primals_7,
            buf177, buf193, buf209, buf181, buf182, buf197, buf198, buf213,
            buf214, 16, XBLOCK=16, num_warps=1, num_stages=1)
        buf223 = buf179
        del buf179
        triton_poi_fused_bmm_8[grid(16)](buf171, buf223, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf224 = reinterpret_tensor(buf178, (4, 1, 4), (4, 16, 1), 0)
        del buf178
        triton_poi_fused_bmm_9[grid(16)](buf170, buf224, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf225 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf223, buf224, out=buf225)
        buf184 = reinterpret_tensor(buf224, (4, 4, 1), (4, 1, 16), 0)
        del buf224
        triton_poi_fused_bmm_13[grid(16)](buf171, buf184, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf185 = reinterpret_tensor(buf223, (4, 4, 1), (4, 1, 1), 0)
        del buf223
        extern_kernels.bmm(buf180, buf184, out=buf185)
        buf200 = buf184
        del buf184
        triton_poi_fused_bmm_14[grid(16)](buf171, buf200, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf201 = empty_strided_cuda((4, 4, 1), (4, 1, 1), torch.float32)
        extern_kernels.bmm(buf196, buf200, out=buf201)
        buf216 = buf200
        del buf200
        triton_poi_fused_bmm_15[grid(16)](buf171, buf216, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf217 = empty_strided_cuda((4, 4, 1), (4, 1, 1), torch.float32)
        extern_kernels.bmm(buf212, buf216, out=buf217)
        buf232 = buf216
        del buf216
        triton_poi_fused_bmm_16[grid(16)](buf171, buf232, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf233 = empty_strided_cuda((4, 4, 1), (4, 1, 1), torch.float32)
        extern_kernels.bmm(buf228, buf232, out=buf233)
        buf239 = empty_strided_cuda((4, 4, 8), (32, 8, 1), torch.float32)
        buf234 = reinterpret_tensor(buf239, (4, 4, 4), (32, 8, 1), 4)
        triton_poi_fused_cat_17[grid(64)](buf185, buf201, buf217, buf233,
            buf234, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf238 = reinterpret_tensor(buf239, (4, 4, 4), (32, 8, 1), 0)
        triton_poi_fused_cat_18[grid(64)](buf163, buf238, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf242 = buf163
        del buf163
        extern_kernels.mm(reinterpret_tensor(buf239, (16, 8), (8, 1), 0),
            reinterpret_tensor(primals_13, (8, 4), (1, 8), 0), out=buf242)
        buf229 = reinterpret_tensor(buf233, (4, 4, 1), (4, 1, 16), 0)
        del buf233
        buf230 = reinterpret_tensor(buf217, (4, 4, 1), (4, 1, 16), 0)
        del buf217
        buf244 = reinterpret_tensor(buf201, (4, 4), (4, 1), 0)
        del buf201
        buf246 = buf244
        del buf244
        triton_poi_fused__softmax_div_eq_masked_fill_mul_relu_sum_19[grid(16)](
            buf246, primals_7, buf225, buf242, primals_14, buf229, buf230, 
            16, XBLOCK=16, num_warps=1, num_stages=1)
        buf183 = buf177
        del buf177
        buf199 = buf193
        del buf193
        buf215 = buf209
        del buf209
        buf231 = buf225
        del buf225
        triton_poi_fused__softmax_eq_masked_fill_11[grid(64)](buf183,
            buf199, buf215, buf231, primals_7, buf181, buf182, buf197,
            buf198, buf213, buf214, buf229, buf230, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf186 = buf230
        del buf230
        triton_poi_fused_bmm_13[grid(16)](buf170, buf186, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf187 = reinterpret_tensor(buf229, (4, 4, 1), (4, 1, 1), 0)
        del buf229
        extern_kernels.bmm(buf183, buf186, out=buf187)
        buf202 = buf186
        del buf186
        triton_poi_fused_bmm_14[grid(16)](buf170, buf202, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf203 = reinterpret_tensor(buf214, (4, 4, 1), (4, 1, 1), 0)
        del buf214
        extern_kernels.bmm(buf199, buf202, out=buf203)
        buf218 = buf202
        del buf202
        triton_poi_fused_bmm_15[grid(16)](buf170, buf218, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf219 = reinterpret_tensor(buf213, (4, 4, 1), (4, 1, 1), 0)
        del buf213
        extern_kernels.bmm(buf215, buf218, out=buf219)
        buf235 = buf218
        del buf218
        triton_poi_fused_bmm_16[grid(16)](buf170, buf235, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf236 = reinterpret_tensor(buf198, (4, 4, 1), (4, 1, 1), 0)
        del buf198
        extern_kernels.bmm(buf231, buf235, out=buf236)
        buf241 = empty_strided_cuda((4, 4, 8), (32, 8, 1), torch.float32)
        buf237 = reinterpret_tensor(buf241, (4, 4, 4), (32, 8, 1), 4)
        triton_poi_fused_cat_17[grid(64)](buf187, buf203, buf219, buf236,
            buf237, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf240 = reinterpret_tensor(buf241, (4, 4, 4), (32, 8, 1), 0)
        triton_poi_fused_cat_18[grid(64)](buf165, buf240, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf243 = buf165
        del buf165
        extern_kernels.mm(reinterpret_tensor(buf241, (16, 8), (8, 1), 0),
            reinterpret_tensor(primals_15, (8, 4), (1, 8), 0), out=buf243)
        buf245 = reinterpret_tensor(buf236, (4, 4), (4, 1), 0)
        del buf236
        buf248 = buf245
        del buf245
        triton_poi_fused_div_mul_relu_sum_20[grid(16)](buf248, buf243,
            primals_16, primals_8, 16, XBLOCK=16, num_warps=1, num_stages=1)
        buf247 = reinterpret_tensor(buf219, (4, 4), (4, 1), 0)
        del buf219
        extern_kernels.addmm(primals_18, buf246, reinterpret_tensor(
            primals_17, (4, 4), (1, 4), 0), alpha=1, beta=1, out=buf247)
        buf249 = reinterpret_tensor(buf203, (4, 4), (4, 1), 0)
        del buf203
        extern_kernels.addmm(primals_20, buf248, reinterpret_tensor(
            primals_19, (4, 4), (1, 4), 0), alpha=1, beta=1, out=buf249)
        buf250 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf667 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.bool)
        triton_poi_fused_relu_threshold_backward_21[grid(64)](buf242,
            primals_14, buf250, buf667, 64, XBLOCK=64, num_warps=1,
            num_stages=1)
        buf251 = empty_strided_cuda((16, 12), (12, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(buf250, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_21, (4, 12), (1, 4), 0), out=buf251)
        buf252 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf666 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.bool)
        triton_poi_fused_relu_threshold_backward_21[grid(64)](buf243,
            primals_16, buf252, buf666, 64, XBLOCK=64, num_warps=1,
            num_stages=1)
        buf253 = empty_strided_cuda((16, 12), (12, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(buf252, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_23, (4, 12), (1, 4), 0), out=buf253)
        buf254 = reinterpret_tensor(buf251, (4, 4, 12), (48, 12, 1), 0)
        del buf251
        triton_poi_fused_mul_1[grid(192)](buf254, primals_22, primals_7, 
            192, XBLOCK=256, num_warps=4, num_stages=1)
        buf255 = reinterpret_tensor(buf253, (4, 4, 12), (48, 12, 1), 0)
        del buf253
        triton_poi_fused_mul_1[grid(192)](buf255, primals_24, primals_8, 
            192, XBLOCK=256, num_warps=4, num_stages=1)
        buf256 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf257 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        triton_poi_fused_add_mul_22[grid(64)](buf249, buf254, buf256,
            buf257, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf258 = reinterpret_tensor(buf187, (4, 4, 1), (4, 1, 16), 0)
        del buf187
        triton_poi_fused_bmm_23[grid(16)](buf256, buf258, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf259 = reinterpret_tensor(buf235, (4, 1, 4), (4, 16, 1), 0)
        del buf235
        triton_poi_fused_bmm_23[grid(16)](buf257, buf259, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf260 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf258, buf259, out=buf260)
        buf261 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf262 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        triton_poi_fused_add_mul_22[grid(64)](buf247, buf255, buf261,
            buf262, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf263 = reinterpret_tensor(buf259, (4, 4, 1), (4, 1, 16), 0)
        del buf259
        triton_poi_fused_bmm_23[grid(16)](buf261, buf263, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf264 = reinterpret_tensor(buf258, (4, 1, 4), (4, 16, 1), 0)
        del buf258
        triton_poi_fused_bmm_23[grid(16)](buf262, buf264, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf265 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf263, buf264, out=buf265)
        buf276 = reinterpret_tensor(buf264, (4, 4, 1), (4, 1, 16), 0)
        del buf264
        triton_poi_fused_bmm_24[grid(16)](buf256, buf276, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf277 = reinterpret_tensor(buf263, (4, 1, 4), (4, 16, 1), 0)
        del buf263
        triton_poi_fused_bmm_24[grid(16)](buf257, buf277, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf278 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf276, buf277, out=buf278)
        buf292 = reinterpret_tensor(buf277, (4, 4, 1), (4, 1, 16), 0)
        del buf277
        triton_poi_fused_bmm_25[grid(16)](buf256, buf292, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf293 = reinterpret_tensor(buf276, (4, 1, 4), (4, 16, 1), 0)
        del buf276
        triton_poi_fused_bmm_25[grid(16)](buf257, buf293, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf294 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf292, buf293, out=buf294)
        buf308 = reinterpret_tensor(buf293, (4, 4, 1), (4, 1, 16), 0)
        del buf293
        triton_poi_fused_bmm_26[grid(16)](buf256, buf308, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf309 = reinterpret_tensor(buf292, (4, 1, 4), (4, 16, 1), 0)
        del buf292
        triton_poi_fused_bmm_26[grid(16)](buf257, buf309, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf310 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf308, buf309, out=buf310)
        buf266 = reinterpret_tensor(buf309, (4, 4, 1), (4, 1, 16), 0)
        del buf309
        buf267 = buf308
        del buf308
        buf282 = buf197
        del buf197
        buf283 = buf182
        del buf182
        buf298 = buf181
        del buf181
        buf299 = reinterpret_tensor(buf185, (4, 4, 1), (4, 1, 16), 0)
        del buf185
        buf314 = buf232
        del buf232
        buf315 = empty_strided_cuda((4, 4, 1), (4, 1, 16), torch.float32)
        triton_poi_fused__softmax_eq_masked_fill_10[grid(16)](primals_7,
            buf260, buf278, buf294, buf310, buf266, buf267, buf282, buf283,
            buf298, buf299, buf314, buf315, 16, XBLOCK=16, num_warps=1,
            num_stages=1)
        buf268 = buf260
        del buf260
        buf284 = buf278
        del buf278
        buf300 = buf294
        del buf294
        buf316 = buf310
        del buf310
        triton_poi_fused__softmax_eq_masked_fill_11[grid(64)](buf268,
            buf284, buf300, buf316, primals_7, buf266, buf267, buf282,
            buf283, buf298, buf299, buf314, buf315, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf279 = buf315
        del buf315
        triton_poi_fused_bmm_24[grid(16)](buf261, buf279, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf280 = reinterpret_tensor(buf314, (4, 1, 4), (4, 16, 1), 0)
        del buf314
        triton_poi_fused_bmm_24[grid(16)](buf262, buf280, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf281 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf279, buf280, out=buf281)
        buf295 = reinterpret_tensor(buf280, (4, 4, 1), (4, 1, 16), 0)
        del buf280
        triton_poi_fused_bmm_25[grid(16)](buf261, buf295, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf296 = reinterpret_tensor(buf279, (4, 1, 4), (4, 16, 1), 0)
        del buf279
        triton_poi_fused_bmm_25[grid(16)](buf262, buf296, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf297 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf295, buf296, out=buf297)
        buf311 = reinterpret_tensor(buf296, (4, 4, 1), (4, 1, 16), 0)
        del buf296
        triton_poi_fused_bmm_26[grid(16)](buf261, buf311, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf312 = reinterpret_tensor(buf295, (4, 1, 4), (4, 16, 1), 0)
        del buf295
        triton_poi_fused_bmm_26[grid(16)](buf262, buf312, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf313 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf311, buf312, out=buf313)
        buf269 = reinterpret_tensor(buf312, (4, 4, 1), (4, 1, 16), 0)
        del buf312
        buf270 = buf311
        del buf311
        buf285 = buf299
        del buf299
        buf286 = buf298
        del buf298
        buf301 = buf283
        del buf283
        buf302 = buf282
        del buf282
        buf317 = buf267
        del buf267
        buf318 = buf266
        del buf266
        triton_poi_fused__softmax_eq_masked_fill_10[grid(16)](primals_8,
            buf265, buf281, buf297, buf313, buf269, buf270, buf285, buf286,
            buf301, buf302, buf317, buf318, 16, XBLOCK=16, num_warps=1,
            num_stages=1)
        buf271 = buf265
        del buf265
        buf287 = buf281
        del buf281
        buf303 = buf297
        del buf297
        buf319 = buf313
        del buf313
        triton_poi_fused__softmax_eq_masked_fill_11[grid(64)](buf271,
            buf287, buf303, buf319, primals_8, buf269, buf270, buf285,
            buf286, buf301, buf302, buf317, buf318, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf272 = buf318
        del buf318
        triton_poi_fused_bmm_13[grid(16)](buf254, buf272, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf273 = reinterpret_tensor(buf317, (4, 4, 1), (4, 1, 1), 0)
        del buf317
        extern_kernels.bmm(buf268, buf272, out=buf273)
        buf274 = buf272
        del buf272
        triton_poi_fused_bmm_13[grid(16)](buf255, buf274, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf275 = reinterpret_tensor(buf302, (4, 4, 1), (4, 1, 1), 0)
        del buf302
        extern_kernels.bmm(buf271, buf274, out=buf275)
        buf288 = buf274
        del buf274
        triton_poi_fused_bmm_14[grid(16)](buf254, buf288, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf289 = reinterpret_tensor(buf301, (4, 4, 1), (4, 1, 1), 0)
        del buf301
        extern_kernels.bmm(buf284, buf288, out=buf289)
        buf290 = buf288
        del buf288
        triton_poi_fused_bmm_14[grid(16)](buf255, buf290, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf291 = reinterpret_tensor(buf286, (4, 4, 1), (4, 1, 1), 0)
        del buf286
        extern_kernels.bmm(buf287, buf290, out=buf291)
        buf304 = buf290
        del buf290
        triton_poi_fused_bmm_15[grid(16)](buf254, buf304, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf305 = reinterpret_tensor(buf285, (4, 4, 1), (4, 1, 1), 0)
        del buf285
        extern_kernels.bmm(buf300, buf304, out=buf305)
        buf306 = buf304
        del buf304
        triton_poi_fused_bmm_15[grid(16)](buf255, buf306, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf307 = reinterpret_tensor(buf270, (4, 4, 1), (4, 1, 1), 0)
        del buf270
        extern_kernels.bmm(buf303, buf306, out=buf307)
        buf320 = buf306
        del buf306
        triton_poi_fused_bmm_16[grid(16)](buf254, buf320, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf321 = reinterpret_tensor(buf269, (4, 4, 1), (4, 1, 1), 0)
        del buf269
        extern_kernels.bmm(buf316, buf320, out=buf321)
        buf326 = reinterpret_tensor(buf242, (4, 4, 4), (16, 4, 1), 0)
        del buf242
        triton_poi_fused_add_cat_28[grid(64)](buf326, buf273, buf289,
            buf305, buf321, primals_14, 64, XBLOCK=64, num_warps=1,
            num_stages=1)
        buf323 = reinterpret_tensor(buf321, (4, 4, 1), (4, 1, 16), 0)
        del buf321
        triton_poi_fused_bmm_16[grid(16)](buf255, buf323, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf324 = buf305
        del buf305
        extern_kernels.bmm(buf319, buf323, out=buf324)
        buf328 = reinterpret_tensor(buf243, (4, 4, 4), (16, 4, 1), 0)
        del buf243
        triton_poi_fused_add_cat_28[grid(64)](buf328, buf275, buf291,
            buf307, buf324, primals_16, 64, XBLOCK=64, num_warps=1,
            num_stages=1)
        buf327 = empty_strided_cuda((16, 4), (4, 1), torch.float32)
        extern_kernels.addmm(primals_26, reinterpret_tensor(buf326, (16, 4),
            (4, 1), 0), reinterpret_tensor(primals_25, (4, 4), (1, 4), 0),
            alpha=1, beta=1, out=buf327)
        buf329 = empty_strided_cuda((16, 4), (4, 1), torch.float32)
        extern_kernels.addmm(primals_28, reinterpret_tensor(buf328, (16, 4),
            (4, 1), 0), reinterpret_tensor(primals_27, (4, 4), (1, 4), 0),
            alpha=1, beta=1, out=buf329)
        buf330 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf665 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.bool)
        triton_poi_fused_relu_threshold_backward_0[grid(64)](buf327, buf330,
            buf665, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf331 = empty_strided_cuda((16, 12), (12, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(buf330, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_9, (4, 12), (1, 4), 0), out=buf331)
        buf332 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf664 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.bool)
        triton_poi_fused_relu_threshold_backward_0[grid(64)](buf329, buf332,
            buf664, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf333 = empty_strided_cuda((16, 12), (12, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(buf332, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_11, (4, 12), (1, 4), 0), out=buf333)
        buf334 = reinterpret_tensor(buf331, (4, 4, 12), (48, 12, 1), 0)
        del buf331
        triton_poi_fused_mul_1[grid(192)](buf334, primals_10, primals_7, 
            192, XBLOCK=256, num_warps=4, num_stages=1)
        buf335 = reinterpret_tensor(buf333, (4, 4, 12), (48, 12, 1), 0)
        del buf333
        triton_poi_fused_mul_1[grid(192)](buf335, primals_12, primals_8, 
            192, XBLOCK=256, num_warps=4, num_stages=1)
        buf336 = reinterpret_tensor(buf324, (4, 4, 1), (4, 1, 16), 0)
        del buf324
        triton_poi_fused_bmm_2[grid(16)](buf334, buf336, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf337 = reinterpret_tensor(buf307, (4, 1, 4), (4, 16, 1), 0)
        del buf307
        triton_poi_fused_bmm_3[grid(16)](buf335, buf337, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf338 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf336, buf337, out=buf338)
        buf339 = reinterpret_tensor(buf337, (4, 4, 1), (4, 1, 16), 0)
        del buf337
        triton_poi_fused_bmm_2[grid(16)](buf335, buf339, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf340 = reinterpret_tensor(buf336, (4, 1, 4), (4, 16, 1), 0)
        del buf336
        triton_poi_fused_bmm_3[grid(16)](buf334, buf340, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf341 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf339, buf340, out=buf341)
        buf352 = reinterpret_tensor(buf340, (4, 4, 1), (4, 1, 16), 0)
        del buf340
        triton_poi_fused_bmm_4[grid(16)](buf334, buf352, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf353 = reinterpret_tensor(buf339, (4, 1, 4), (4, 16, 1), 0)
        del buf339
        triton_poi_fused_bmm_5[grid(16)](buf335, buf353, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf354 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf352, buf353, out=buf354)
        buf368 = reinterpret_tensor(buf353, (4, 4, 1), (4, 1, 16), 0)
        del buf353
        triton_poi_fused_bmm_6[grid(16)](buf334, buf368, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf369 = reinterpret_tensor(buf352, (4, 1, 4), (4, 16, 1), 0)
        del buf352
        triton_poi_fused_bmm_7[grid(16)](buf335, buf369, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf370 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf368, buf369, out=buf370)
        buf384 = reinterpret_tensor(buf369, (4, 4, 1), (4, 1, 16), 0)
        del buf369
        triton_poi_fused_bmm_8[grid(16)](buf334, buf384, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf385 = reinterpret_tensor(buf368, (4, 1, 4), (4, 16, 1), 0)
        del buf368
        triton_poi_fused_bmm_9[grid(16)](buf335, buf385, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf386 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf384, buf385, out=buf386)
        buf342 = reinterpret_tensor(buf385, (4, 4, 1), (4, 1, 16), 0)
        del buf385
        buf343 = buf384
        del buf384
        buf358 = reinterpret_tensor(buf291, (4, 4, 1), (4, 1, 16), 0)
        del buf291
        buf359 = reinterpret_tensor(buf275, (4, 4, 1), (4, 1, 16), 0)
        del buf275
        buf374 = buf323
        del buf323
        buf375 = reinterpret_tensor(buf289, (4, 4, 1), (4, 1, 16), 0)
        del buf289
        buf390 = reinterpret_tensor(buf273, (4, 4, 1), (4, 1, 16), 0)
        del buf273
        buf391 = buf320
        del buf320
        triton_poi_fused__softmax_eq_masked_fill_10[grid(16)](primals_8,
            buf338, buf354, buf370, buf386, buf342, buf343, buf358, buf359,
            buf374, buf375, buf390, buf391, 16, XBLOCK=16, num_warps=1,
            num_stages=1)
        buf344 = buf338
        del buf338
        buf360 = buf354
        del buf354
        buf376 = buf370
        del buf370
        buf392 = buf386
        del buf386
        triton_poi_fused__softmax_eq_masked_fill_11[grid(64)](buf344,
            buf360, buf376, buf392, primals_8, buf342, buf343, buf358,
            buf359, buf374, buf375, buf390, buf391, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf355 = buf391
        del buf391
        triton_poi_fused_bmm_4[grid(16)](buf335, buf355, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf356 = reinterpret_tensor(buf390, (4, 1, 4), (4, 16, 1), 0)
        del buf390
        triton_poi_fused_bmm_5[grid(16)](buf334, buf356, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf357 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf355, buf356, out=buf357)
        buf371 = reinterpret_tensor(buf356, (4, 4, 1), (4, 1, 16), 0)
        del buf356
        triton_poi_fused_bmm_6[grid(16)](buf335, buf371, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf372 = reinterpret_tensor(buf355, (4, 1, 4), (4, 16, 1), 0)
        del buf355
        triton_poi_fused_bmm_7[grid(16)](buf334, buf372, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf373 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf371, buf372, out=buf373)
        buf345 = reinterpret_tensor(buf372, (4, 4, 1), (4, 1, 16), 0)
        del buf372
        buf346 = buf371
        del buf371
        buf361 = buf375
        del buf375
        buf362 = buf374
        del buf374
        buf377 = buf359
        del buf359
        buf378 = buf358
        del buf358
        triton_poi_fused__softmax_eq_masked_fill_12[grid(16)](primals_7,
            buf341, buf357, buf373, buf345, buf346, buf361, buf362, buf377,
            buf378, 16, XBLOCK=16, num_warps=1, num_stages=1)
        buf387 = buf343
        del buf343
        triton_poi_fused_bmm_8[grid(16)](buf335, buf387, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf388 = reinterpret_tensor(buf342, (4, 1, 4), (4, 16, 1), 0)
        del buf342
        triton_poi_fused_bmm_9[grid(16)](buf334, buf388, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf389 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf387, buf388, out=buf389)
        buf348 = reinterpret_tensor(buf388, (4, 4, 1), (4, 1, 16), 0)
        del buf388
        triton_poi_fused_bmm_13[grid(16)](buf335, buf348, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf349 = reinterpret_tensor(buf387, (4, 4, 1), (4, 1, 1), 0)
        del buf387
        extern_kernels.bmm(buf344, buf348, out=buf349)
        buf364 = buf348
        del buf348
        triton_poi_fused_bmm_14[grid(16)](buf335, buf364, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf365 = empty_strided_cuda((4, 4, 1), (4, 1, 1), torch.float32)
        extern_kernels.bmm(buf360, buf364, out=buf365)
        buf380 = buf364
        del buf364
        triton_poi_fused_bmm_15[grid(16)](buf335, buf380, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf381 = empty_strided_cuda((4, 4, 1), (4, 1, 1), torch.float32)
        extern_kernels.bmm(buf376, buf380, out=buf381)
        buf396 = buf380
        del buf380
        triton_poi_fused_bmm_16[grid(16)](buf335, buf396, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf397 = empty_strided_cuda((4, 4, 1), (4, 1, 1), torch.float32)
        extern_kernels.bmm(buf392, buf396, out=buf397)
        buf403 = empty_strided_cuda((4, 4, 8), (32, 8, 1), torch.float32)
        buf398 = reinterpret_tensor(buf403, (4, 4, 4), (32, 8, 1), 4)
        triton_poi_fused_cat_17[grid(64)](buf349, buf365, buf381, buf397,
            buf398, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf402 = reinterpret_tensor(buf403, (4, 4, 4), (32, 8, 1), 0)
        triton_poi_fused_cat_18[grid(64)](buf327, buf402, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf406 = buf327
        del buf327
        extern_kernels.mm(reinterpret_tensor(buf403, (16, 8), (8, 1), 0),
            reinterpret_tensor(primals_13, (8, 4), (1, 8), 0), out=buf406)
        buf393 = reinterpret_tensor(buf397, (4, 4, 1), (4, 1, 16), 0)
        del buf397
        buf394 = reinterpret_tensor(buf381, (4, 4, 1), (4, 1, 16), 0)
        del buf381
        buf408 = reinterpret_tensor(buf365, (4, 4), (4, 1), 0)
        del buf365
        buf410 = buf408
        del buf408
        triton_poi_fused__softmax_div_eq_masked_fill_mul_relu_sum_19[grid(16)](
            buf410, primals_7, buf389, buf406, primals_14, buf393, buf394, 
            16, XBLOCK=16, num_warps=1, num_stages=1)
        buf347 = buf341
        del buf341
        buf363 = buf357
        del buf357
        buf379 = buf373
        del buf373
        buf395 = buf389
        del buf389
        triton_poi_fused__softmax_eq_masked_fill_11[grid(64)](buf347,
            buf363, buf379, buf395, primals_7, buf345, buf346, buf361,
            buf362, buf377, buf378, buf393, buf394, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf350 = buf394
        del buf394
        triton_poi_fused_bmm_13[grid(16)](buf334, buf350, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf351 = reinterpret_tensor(buf393, (4, 4, 1), (4, 1, 1), 0)
        del buf393
        extern_kernels.bmm(buf347, buf350, out=buf351)
        buf366 = buf350
        del buf350
        triton_poi_fused_bmm_14[grid(16)](buf334, buf366, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf367 = reinterpret_tensor(buf378, (4, 4, 1), (4, 1, 1), 0)
        del buf378
        extern_kernels.bmm(buf363, buf366, out=buf367)
        buf382 = buf366
        del buf366
        triton_poi_fused_bmm_15[grid(16)](buf334, buf382, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf383 = reinterpret_tensor(buf377, (4, 4, 1), (4, 1, 1), 0)
        del buf377
        extern_kernels.bmm(buf379, buf382, out=buf383)
        buf399 = buf382
        del buf382
        triton_poi_fused_bmm_16[grid(16)](buf334, buf399, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf400 = reinterpret_tensor(buf362, (4, 4, 1), (4, 1, 1), 0)
        del buf362
        extern_kernels.bmm(buf395, buf399, out=buf400)
        buf405 = empty_strided_cuda((4, 4, 8), (32, 8, 1), torch.float32)
        buf401 = reinterpret_tensor(buf405, (4, 4, 4), (32, 8, 1), 4)
        triton_poi_fused_cat_17[grid(64)](buf351, buf367, buf383, buf400,
            buf401, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf404 = reinterpret_tensor(buf405, (4, 4, 4), (32, 8, 1), 0)
        triton_poi_fused_cat_18[grid(64)](buf329, buf404, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf407 = buf329
        del buf329
        extern_kernels.mm(reinterpret_tensor(buf405, (16, 8), (8, 1), 0),
            reinterpret_tensor(primals_15, (8, 4), (1, 8), 0), out=buf407)
        buf409 = reinterpret_tensor(buf400, (4, 4), (4, 1), 0)
        del buf400
        buf412 = buf409
        del buf409
        triton_poi_fused_div_mul_relu_sum_20[grid(16)](buf412, buf407,
            primals_16, primals_8, 16, XBLOCK=16, num_warps=1, num_stages=1)
        buf411 = reinterpret_tensor(buf383, (4, 4), (4, 1), 0)
        del buf383
        extern_kernels.addmm(primals_18, buf410, reinterpret_tensor(
            primals_17, (4, 4), (1, 4), 0), alpha=1, beta=1, out=buf411)
        buf413 = reinterpret_tensor(buf367, (4, 4), (4, 1), 0)
        del buf367
        extern_kernels.addmm(primals_20, buf412, reinterpret_tensor(
            primals_19, (4, 4), (1, 4), 0), alpha=1, beta=1, out=buf413)
        buf414 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf663 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.bool)
        triton_poi_fused_relu_threshold_backward_21[grid(64)](buf406,
            primals_14, buf414, buf663, 64, XBLOCK=64, num_warps=1,
            num_stages=1)
        buf415 = empty_strided_cuda((16, 12), (12, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(buf414, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_21, (4, 12), (1, 4), 0), out=buf415)
        buf416 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf662 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.bool)
        triton_poi_fused_relu_threshold_backward_21[grid(64)](buf407,
            primals_16, buf416, buf662, 64, XBLOCK=64, num_warps=1,
            num_stages=1)
        buf417 = empty_strided_cuda((16, 12), (12, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(buf416, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_23, (4, 12), (1, 4), 0), out=buf417)
        buf418 = reinterpret_tensor(buf415, (4, 4, 12), (48, 12, 1), 0)
        del buf415
        triton_poi_fused_mul_1[grid(192)](buf418, primals_22, primals_7, 
            192, XBLOCK=256, num_warps=4, num_stages=1)
        buf419 = reinterpret_tensor(buf417, (4, 4, 12), (48, 12, 1), 0)
        del buf417
        triton_poi_fused_mul_1[grid(192)](buf419, primals_24, primals_8, 
            192, XBLOCK=256, num_warps=4, num_stages=1)
        buf420 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf421 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        triton_poi_fused_add_mul_22[grid(64)](buf413, buf418, buf420,
            buf421, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf422 = reinterpret_tensor(buf351, (4, 4, 1), (4, 1, 16), 0)
        del buf351
        triton_poi_fused_bmm_23[grid(16)](buf420, buf422, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf423 = reinterpret_tensor(buf399, (4, 1, 4), (4, 16, 1), 0)
        del buf399
        triton_poi_fused_bmm_23[grid(16)](buf421, buf423, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf424 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf422, buf423, out=buf424)
        buf425 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf426 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        triton_poi_fused_add_mul_22[grid(64)](buf411, buf419, buf425,
            buf426, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf427 = reinterpret_tensor(buf423, (4, 4, 1), (4, 1, 16), 0)
        del buf423
        triton_poi_fused_bmm_23[grid(16)](buf425, buf427, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf428 = reinterpret_tensor(buf422, (4, 1, 4), (4, 16, 1), 0)
        del buf422
        triton_poi_fused_bmm_23[grid(16)](buf426, buf428, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf429 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf427, buf428, out=buf429)
        buf440 = reinterpret_tensor(buf428, (4, 4, 1), (4, 1, 16), 0)
        del buf428
        triton_poi_fused_bmm_24[grid(16)](buf420, buf440, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf441 = reinterpret_tensor(buf427, (4, 1, 4), (4, 16, 1), 0)
        del buf427
        triton_poi_fused_bmm_24[grid(16)](buf421, buf441, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf442 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf440, buf441, out=buf442)
        buf456 = reinterpret_tensor(buf441, (4, 4, 1), (4, 1, 16), 0)
        del buf441
        triton_poi_fused_bmm_25[grid(16)](buf420, buf456, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf457 = reinterpret_tensor(buf440, (4, 1, 4), (4, 16, 1), 0)
        del buf440
        triton_poi_fused_bmm_25[grid(16)](buf421, buf457, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf458 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf456, buf457, out=buf458)
        buf472 = reinterpret_tensor(buf457, (4, 4, 1), (4, 1, 16), 0)
        del buf457
        triton_poi_fused_bmm_26[grid(16)](buf420, buf472, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf473 = reinterpret_tensor(buf456, (4, 1, 4), (4, 16, 1), 0)
        del buf456
        triton_poi_fused_bmm_26[grid(16)](buf421, buf473, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf474 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf472, buf473, out=buf474)
        buf430 = reinterpret_tensor(buf473, (4, 4, 1), (4, 1, 16), 0)
        del buf473
        buf431 = buf472
        del buf472
        buf446 = buf361
        del buf361
        buf447 = buf346
        del buf346
        buf462 = buf345
        del buf345
        buf463 = reinterpret_tensor(buf349, (4, 4, 1), (4, 1, 16), 0)
        del buf349
        buf478 = buf396
        del buf396
        buf479 = empty_strided_cuda((4, 4, 1), (4, 1, 16), torch.float32)
        triton_poi_fused__softmax_eq_masked_fill_10[grid(16)](primals_7,
            buf424, buf442, buf458, buf474, buf430, buf431, buf446, buf447,
            buf462, buf463, buf478, buf479, 16, XBLOCK=16, num_warps=1,
            num_stages=1)
        buf432 = buf424
        del buf424
        buf448 = buf442
        del buf442
        buf464 = buf458
        del buf458
        buf480 = buf474
        del buf474
        triton_poi_fused__softmax_eq_masked_fill_11[grid(64)](buf432,
            buf448, buf464, buf480, primals_7, buf430, buf431, buf446,
            buf447, buf462, buf463, buf478, buf479, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf443 = buf479
        del buf479
        triton_poi_fused_bmm_24[grid(16)](buf425, buf443, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf444 = reinterpret_tensor(buf478, (4, 1, 4), (4, 16, 1), 0)
        del buf478
        triton_poi_fused_bmm_24[grid(16)](buf426, buf444, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf445 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf443, buf444, out=buf445)
        buf459 = reinterpret_tensor(buf444, (4, 4, 1), (4, 1, 16), 0)
        del buf444
        triton_poi_fused_bmm_25[grid(16)](buf425, buf459, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf460 = reinterpret_tensor(buf443, (4, 1, 4), (4, 16, 1), 0)
        del buf443
        triton_poi_fused_bmm_25[grid(16)](buf426, buf460, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf461 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf459, buf460, out=buf461)
        buf475 = reinterpret_tensor(buf460, (4, 4, 1), (4, 1, 16), 0)
        del buf460
        triton_poi_fused_bmm_26[grid(16)](buf425, buf475, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf476 = reinterpret_tensor(buf459, (4, 1, 4), (4, 16, 1), 0)
        del buf459
        triton_poi_fused_bmm_26[grid(16)](buf426, buf476, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf477 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf475, buf476, out=buf477)
        buf433 = reinterpret_tensor(buf476, (4, 4, 1), (4, 1, 16), 0)
        del buf476
        buf434 = buf475
        del buf475
        buf449 = buf463
        del buf463
        buf450 = buf462
        del buf462
        buf465 = buf447
        del buf447
        buf466 = buf446
        del buf446
        buf481 = buf431
        del buf431
        buf482 = buf430
        del buf430
        triton_poi_fused__softmax_eq_masked_fill_10[grid(16)](primals_8,
            buf429, buf445, buf461, buf477, buf433, buf434, buf449, buf450,
            buf465, buf466, buf481, buf482, 16, XBLOCK=16, num_warps=1,
            num_stages=1)
        buf435 = buf429
        del buf429
        buf451 = buf445
        del buf445
        buf467 = buf461
        del buf461
        buf483 = buf477
        del buf477
        triton_poi_fused__softmax_eq_masked_fill_11[grid(64)](buf435,
            buf451, buf467, buf483, primals_8, buf433, buf434, buf449,
            buf450, buf465, buf466, buf481, buf482, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf436 = buf482
        del buf482
        triton_poi_fused_bmm_13[grid(16)](buf418, buf436, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf437 = reinterpret_tensor(buf481, (4, 4, 1), (4, 1, 1), 0)
        del buf481
        extern_kernels.bmm(buf432, buf436, out=buf437)
        buf438 = buf436
        del buf436
        triton_poi_fused_bmm_13[grid(16)](buf419, buf438, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf439 = reinterpret_tensor(buf466, (4, 4, 1), (4, 1, 1), 0)
        del buf466
        extern_kernels.bmm(buf435, buf438, out=buf439)
        buf452 = buf438
        del buf438
        triton_poi_fused_bmm_14[grid(16)](buf418, buf452, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf453 = reinterpret_tensor(buf465, (4, 4, 1), (4, 1, 1), 0)
        del buf465
        extern_kernels.bmm(buf448, buf452, out=buf453)
        buf454 = buf452
        del buf452
        triton_poi_fused_bmm_14[grid(16)](buf419, buf454, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf455 = reinterpret_tensor(buf450, (4, 4, 1), (4, 1, 1), 0)
        del buf450
        extern_kernels.bmm(buf451, buf454, out=buf455)
        buf468 = buf454
        del buf454
        triton_poi_fused_bmm_15[grid(16)](buf418, buf468, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf469 = reinterpret_tensor(buf449, (4, 4, 1), (4, 1, 1), 0)
        del buf449
        extern_kernels.bmm(buf464, buf468, out=buf469)
        buf470 = buf468
        del buf468
        triton_poi_fused_bmm_15[grid(16)](buf419, buf470, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf471 = reinterpret_tensor(buf434, (4, 4, 1), (4, 1, 1), 0)
        del buf434
        extern_kernels.bmm(buf467, buf470, out=buf471)
        buf484 = buf470
        del buf470
        triton_poi_fused_bmm_16[grid(16)](buf418, buf484, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf485 = reinterpret_tensor(buf433, (4, 4, 1), (4, 1, 1), 0)
        del buf433
        extern_kernels.bmm(buf480, buf484, out=buf485)
        buf490 = reinterpret_tensor(buf406, (4, 4, 4), (16, 4, 1), 0)
        del buf406
        triton_poi_fused_add_cat_28[grid(64)](buf490, buf437, buf453,
            buf469, buf485, primals_14, 64, XBLOCK=64, num_warps=1,
            num_stages=1)
        buf487 = reinterpret_tensor(buf485, (4, 4, 1), (4, 1, 16), 0)
        del buf485
        triton_poi_fused_bmm_16[grid(16)](buf419, buf487, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf488 = buf469
        del buf469
        extern_kernels.bmm(buf483, buf487, out=buf488)
        buf492 = reinterpret_tensor(buf407, (4, 4, 4), (16, 4, 1), 0)
        del buf407
        triton_poi_fused_add_cat_28[grid(64)](buf492, buf439, buf455,
            buf471, buf488, primals_16, 64, XBLOCK=64, num_warps=1,
            num_stages=1)
        buf491 = empty_strided_cuda((16, 4), (4, 1), torch.float32)
        extern_kernels.addmm(primals_26, reinterpret_tensor(buf490, (16, 4),
            (4, 1), 0), reinterpret_tensor(primals_25, (4, 4), (1, 4), 0),
            alpha=1, beta=1, out=buf491)
        buf493 = empty_strided_cuda((16, 4), (4, 1), torch.float32)
        extern_kernels.addmm(primals_28, reinterpret_tensor(buf492, (16, 4),
            (4, 1), 0), reinterpret_tensor(primals_27, (4, 4), (1, 4), 0),
            alpha=1, beta=1, out=buf493)
        buf494 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf661 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.bool)
        triton_poi_fused_relu_threshold_backward_0[grid(64)](buf491, buf494,
            buf661, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf495 = empty_strided_cuda((16, 12), (12, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(buf494, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_9, (4, 12), (1, 4), 0), out=buf495)
        buf496 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf660 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.bool)
        triton_poi_fused_relu_threshold_backward_0[grid(64)](buf493, buf496,
            buf660, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf497 = empty_strided_cuda((16, 12), (12, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(buf496, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_11, (4, 12), (1, 4), 0), out=buf497)
        buf498 = reinterpret_tensor(buf495, (4, 4, 12), (48, 12, 1), 0)
        del buf495
        triton_poi_fused_mul_1[grid(192)](buf498, primals_10, primals_7, 
            192, XBLOCK=256, num_warps=4, num_stages=1)
        del primals_10
        buf499 = reinterpret_tensor(buf497, (4, 4, 12), (48, 12, 1), 0)
        del buf497
        triton_poi_fused_mul_1[grid(192)](buf499, primals_12, primals_8, 
            192, XBLOCK=256, num_warps=4, num_stages=1)
        del primals_12
        buf500 = reinterpret_tensor(buf488, (4, 4, 1), (4, 1, 16), 0)
        del buf488
        triton_poi_fused_bmm_2[grid(16)](buf498, buf500, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf501 = reinterpret_tensor(buf471, (4, 1, 4), (4, 16, 1), 0)
        del buf471
        triton_poi_fused_bmm_3[grid(16)](buf499, buf501, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf502 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf500, buf501, out=buf502)
        buf503 = reinterpret_tensor(buf501, (4, 4, 1), (4, 1, 16), 0)
        del buf501
        triton_poi_fused_bmm_2[grid(16)](buf499, buf503, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf504 = reinterpret_tensor(buf500, (4, 1, 4), (4, 16, 1), 0)
        del buf500
        triton_poi_fused_bmm_3[grid(16)](buf498, buf504, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf505 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf503, buf504, out=buf505)
        buf516 = reinterpret_tensor(buf504, (4, 4, 1), (4, 1, 16), 0)
        del buf504
        triton_poi_fused_bmm_4[grid(16)](buf498, buf516, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf517 = reinterpret_tensor(buf503, (4, 1, 4), (4, 16, 1), 0)
        del buf503
        triton_poi_fused_bmm_5[grid(16)](buf499, buf517, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf518 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf516, buf517, out=buf518)
        buf532 = reinterpret_tensor(buf517, (4, 4, 1), (4, 1, 16), 0)
        del buf517
        triton_poi_fused_bmm_6[grid(16)](buf498, buf532, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf533 = reinterpret_tensor(buf516, (4, 1, 4), (4, 16, 1), 0)
        del buf516
        triton_poi_fused_bmm_7[grid(16)](buf499, buf533, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf534 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf532, buf533, out=buf534)
        buf548 = reinterpret_tensor(buf533, (4, 4, 1), (4, 1, 16), 0)
        del buf533
        triton_poi_fused_bmm_8[grid(16)](buf498, buf548, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf549 = reinterpret_tensor(buf532, (4, 1, 4), (4, 16, 1), 0)
        del buf532
        triton_poi_fused_bmm_9[grid(16)](buf499, buf549, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf550 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf548, buf549, out=buf550)
        buf506 = reinterpret_tensor(buf549, (4, 4, 1), (4, 1, 16), 0)
        del buf549
        buf507 = buf548
        del buf548
        buf522 = reinterpret_tensor(buf455, (4, 4, 1), (4, 1, 16), 0)
        del buf455
        buf523 = reinterpret_tensor(buf439, (4, 4, 1), (4, 1, 16), 0)
        del buf439
        buf538 = buf487
        del buf487
        buf539 = reinterpret_tensor(buf453, (4, 4, 1), (4, 1, 16), 0)
        del buf453
        buf554 = reinterpret_tensor(buf437, (4, 4, 1), (4, 1, 16), 0)
        del buf437
        buf555 = buf484
        del buf484
        triton_poi_fused__softmax_eq_masked_fill_10[grid(16)](primals_8,
            buf502, buf518, buf534, buf550, buf506, buf507, buf522, buf523,
            buf538, buf539, buf554, buf555, 16, XBLOCK=16, num_warps=1,
            num_stages=1)
        buf508 = buf502
        del buf502
        buf524 = buf518
        del buf518
        buf540 = buf534
        del buf534
        buf556 = buf550
        del buf550
        triton_poi_fused__softmax_eq_masked_fill_11[grid(64)](buf508,
            buf524, buf540, buf556, primals_8, buf506, buf507, buf522,
            buf523, buf538, buf539, buf554, buf555, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf519 = buf555
        del buf555
        triton_poi_fused_bmm_4[grid(16)](buf499, buf519, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf520 = reinterpret_tensor(buf554, (4, 1, 4), (4, 16, 1), 0)
        del buf554
        triton_poi_fused_bmm_5[grid(16)](buf498, buf520, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf521 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf519, buf520, out=buf521)
        buf535 = reinterpret_tensor(buf520, (4, 4, 1), (4, 1, 16), 0)
        del buf520
        triton_poi_fused_bmm_6[grid(16)](buf499, buf535, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf536 = reinterpret_tensor(buf519, (4, 1, 4), (4, 16, 1), 0)
        del buf519
        triton_poi_fused_bmm_7[grid(16)](buf498, buf536, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf537 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf535, buf536, out=buf537)
        buf509 = reinterpret_tensor(buf536, (4, 4, 1), (4, 1, 16), 0)
        del buf536
        buf510 = buf535
        del buf535
        buf525 = buf539
        del buf539
        buf526 = buf538
        del buf538
        buf541 = buf523
        del buf523
        buf542 = buf522
        del buf522
        triton_poi_fused__softmax_eq_masked_fill_12[grid(16)](primals_7,
            buf505, buf521, buf537, buf509, buf510, buf525, buf526, buf541,
            buf542, 16, XBLOCK=16, num_warps=1, num_stages=1)
        buf551 = buf507
        del buf507
        triton_poi_fused_bmm_8[grid(16)](buf499, buf551, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf552 = reinterpret_tensor(buf506, (4, 1, 4), (4, 16, 1), 0)
        del buf506
        triton_poi_fused_bmm_9[grid(16)](buf498, buf552, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf553 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf551, buf552, out=buf553)
        buf512 = reinterpret_tensor(buf552, (4, 4, 1), (4, 1, 16), 0)
        del buf552
        triton_poi_fused_bmm_13[grid(16)](buf499, buf512, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf513 = reinterpret_tensor(buf551, (4, 4, 1), (4, 1, 1), 0)
        del buf551
        extern_kernels.bmm(buf508, buf512, out=buf513)
        buf528 = buf512
        del buf512
        triton_poi_fused_bmm_14[grid(16)](buf499, buf528, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf529 = empty_strided_cuda((4, 4, 1), (4, 1, 1), torch.float32)
        extern_kernels.bmm(buf524, buf528, out=buf529)
        buf544 = buf528
        del buf528
        triton_poi_fused_bmm_15[grid(16)](buf499, buf544, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf545 = empty_strided_cuda((4, 4, 1), (4, 1, 1), torch.float32)
        extern_kernels.bmm(buf540, buf544, out=buf545)
        buf560 = buf544
        del buf544
        triton_poi_fused_bmm_16[grid(16)](buf499, buf560, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf561 = empty_strided_cuda((4, 4, 1), (4, 1, 1), torch.float32)
        extern_kernels.bmm(buf556, buf560, out=buf561)
        buf567 = empty_strided_cuda((4, 4, 8), (32, 8, 1), torch.float32)
        buf562 = reinterpret_tensor(buf567, (4, 4, 4), (32, 8, 1), 4)
        triton_poi_fused_cat_17[grid(64)](buf513, buf529, buf545, buf561,
            buf562, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf566 = reinterpret_tensor(buf567, (4, 4, 4), (32, 8, 1), 0)
        triton_poi_fused_cat_18[grid(64)](buf491, buf566, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf570 = buf491
        del buf491
        extern_kernels.mm(reinterpret_tensor(buf567, (16, 8), (8, 1), 0),
            reinterpret_tensor(primals_13, (8, 4), (1, 8), 0), out=buf570)
        buf557 = reinterpret_tensor(buf561, (4, 4, 1), (4, 1, 16), 0)
        del buf561
        buf558 = reinterpret_tensor(buf545, (4, 4, 1), (4, 1, 16), 0)
        del buf545
        buf572 = reinterpret_tensor(buf529, (4, 4), (4, 1), 0)
        del buf529
        buf574 = buf572
        del buf572
        triton_poi_fused__softmax_div_eq_masked_fill_mul_relu_sum_19[grid(16)](
            buf574, primals_7, buf553, buf570, primals_14, buf557, buf558, 
            16, XBLOCK=16, num_warps=1, num_stages=1)
        buf511 = buf505
        del buf505
        buf527 = buf521
        del buf521
        buf543 = buf537
        del buf537
        buf559 = buf553
        del buf553
        triton_poi_fused__softmax_eq_masked_fill_11[grid(64)](buf511,
            buf527, buf543, buf559, primals_7, buf509, buf510, buf525,
            buf526, buf541, buf542, buf557, buf558, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf514 = buf558
        del buf558
        triton_poi_fused_bmm_13[grid(16)](buf498, buf514, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf515 = reinterpret_tensor(buf557, (4, 4, 1), (4, 1, 1), 0)
        del buf557
        extern_kernels.bmm(buf511, buf514, out=buf515)
        buf530 = buf514
        del buf514
        triton_poi_fused_bmm_14[grid(16)](buf498, buf530, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf531 = reinterpret_tensor(buf542, (4, 4, 1), (4, 1, 1), 0)
        del buf542
        extern_kernels.bmm(buf527, buf530, out=buf531)
        buf546 = buf530
        del buf530
        triton_poi_fused_bmm_15[grid(16)](buf498, buf546, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf547 = reinterpret_tensor(buf541, (4, 4, 1), (4, 1, 1), 0)
        del buf541
        extern_kernels.bmm(buf543, buf546, out=buf547)
        buf563 = buf546
        del buf546
        triton_poi_fused_bmm_16[grid(16)](buf498, buf563, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf564 = reinterpret_tensor(buf526, (4, 4, 1), (4, 1, 1), 0)
        del buf526
        extern_kernels.bmm(buf559, buf563, out=buf564)
        buf569 = empty_strided_cuda((4, 4, 8), (32, 8, 1), torch.float32)
        buf565 = reinterpret_tensor(buf569, (4, 4, 4), (32, 8, 1), 4)
        triton_poi_fused_cat_17[grid(64)](buf515, buf531, buf547, buf564,
            buf565, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf568 = reinterpret_tensor(buf569, (4, 4, 4), (32, 8, 1), 0)
        triton_poi_fused_cat_18[grid(64)](buf493, buf568, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf571 = buf493
        del buf493
        extern_kernels.mm(reinterpret_tensor(buf569, (16, 8), (8, 1), 0),
            reinterpret_tensor(primals_15, (8, 4), (1, 8), 0), out=buf571)
        buf573 = reinterpret_tensor(buf564, (4, 4), (4, 1), 0)
        del buf564
        buf576 = buf573
        del buf573
        triton_poi_fused_div_mul_relu_sum_20[grid(16)](buf576, buf571,
            primals_16, primals_8, 16, XBLOCK=16, num_warps=1, num_stages=1)
        buf575 = reinterpret_tensor(buf547, (4, 4), (4, 1), 0)
        del buf547
        extern_kernels.addmm(primals_18, buf574, reinterpret_tensor(
            primals_17, (4, 4), (1, 4), 0), alpha=1, beta=1, out=buf575)
        del primals_18
        buf577 = reinterpret_tensor(buf531, (4, 4), (4, 1), 0)
        del buf531
        extern_kernels.addmm(primals_20, buf576, reinterpret_tensor(
            primals_19, (4, 4), (1, 4), 0), alpha=1, beta=1, out=buf577)
        del primals_20
        buf578 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf659 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.bool)
        triton_poi_fused_relu_threshold_backward_21[grid(64)](buf570,
            primals_14, buf578, buf659, 64, XBLOCK=64, num_warps=1,
            num_stages=1)
        buf579 = empty_strided_cuda((16, 12), (12, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(buf578, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_21, (4, 12), (1, 4), 0), out=buf579)
        buf580 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf658 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.bool)
        triton_poi_fused_relu_threshold_backward_21[grid(64)](buf571,
            primals_16, buf580, buf658, 64, XBLOCK=64, num_warps=1,
            num_stages=1)
        buf581 = empty_strided_cuda((16, 12), (12, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(buf580, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_23, (4, 12), (1, 4), 0), out=buf581)
        buf582 = reinterpret_tensor(buf579, (4, 4, 12), (48, 12, 1), 0)
        del buf579
        triton_poi_fused_mul_1[grid(192)](buf582, primals_22, primals_7, 
            192, XBLOCK=256, num_warps=4, num_stages=1)
        del primals_22
        buf583 = reinterpret_tensor(buf581, (4, 4, 12), (48, 12, 1), 0)
        del buf581
        triton_poi_fused_mul_1[grid(192)](buf583, primals_24, primals_8, 
            192, XBLOCK=256, num_warps=4, num_stages=1)
        del primals_24
        buf584 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf585 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        triton_poi_fused_add_mul_22[grid(64)](buf577, buf582, buf584,
            buf585, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf586 = reinterpret_tensor(buf515, (4, 4, 1), (4, 1, 16), 0)
        del buf515
        triton_poi_fused_bmm_23[grid(16)](buf584, buf586, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf587 = reinterpret_tensor(buf563, (4, 1, 4), (4, 16, 1), 0)
        del buf563
        triton_poi_fused_bmm_23[grid(16)](buf585, buf587, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf588 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf586, buf587, out=buf588)
        buf589 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf590 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        triton_poi_fused_add_mul_22[grid(64)](buf575, buf583, buf589,
            buf590, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf591 = reinterpret_tensor(buf587, (4, 4, 1), (4, 1, 16), 0)
        del buf587
        triton_poi_fused_bmm_23[grid(16)](buf589, buf591, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf592 = reinterpret_tensor(buf586, (4, 1, 4), (4, 16, 1), 0)
        del buf586
        triton_poi_fused_bmm_23[grid(16)](buf590, buf592, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf593 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf591, buf592, out=buf593)
        buf604 = reinterpret_tensor(buf592, (4, 4, 1), (4, 1, 16), 0)
        del buf592
        triton_poi_fused_bmm_24[grid(16)](buf584, buf604, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf605 = reinterpret_tensor(buf591, (4, 1, 4), (4, 16, 1), 0)
        del buf591
        triton_poi_fused_bmm_24[grid(16)](buf585, buf605, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf606 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf604, buf605, out=buf606)
        buf620 = reinterpret_tensor(buf605, (4, 4, 1), (4, 1, 16), 0)
        del buf605
        triton_poi_fused_bmm_25[grid(16)](buf584, buf620, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf621 = reinterpret_tensor(buf604, (4, 1, 4), (4, 16, 1), 0)
        del buf604
        triton_poi_fused_bmm_25[grid(16)](buf585, buf621, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf622 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf620, buf621, out=buf622)
        buf636 = reinterpret_tensor(buf621, (4, 4, 1), (4, 1, 16), 0)
        del buf621
        triton_poi_fused_bmm_26[grid(16)](buf584, buf636, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf637 = reinterpret_tensor(buf620, (4, 1, 4), (4, 16, 1), 0)
        del buf620
        triton_poi_fused_bmm_26[grid(16)](buf585, buf637, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf638 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf636, buf637, out=buf638)
        buf594 = reinterpret_tensor(buf637, (4, 4, 1), (4, 1, 16), 0)
        del buf637
        buf595 = buf636
        del buf636
        buf610 = buf525
        del buf525
        buf611 = buf510
        del buf510
        buf626 = buf509
        del buf509
        buf627 = reinterpret_tensor(buf513, (4, 4, 1), (4, 1, 16), 0)
        del buf513
        buf642 = buf560
        del buf560
        buf643 = empty_strided_cuda((4, 4, 1), (4, 1, 16), torch.float32)
        triton_poi_fused__softmax_eq_masked_fill_10[grid(16)](primals_7,
            buf588, buf606, buf622, buf638, buf594, buf595, buf610, buf611,
            buf626, buf627, buf642, buf643, 16, XBLOCK=16, num_warps=1,
            num_stages=1)
        buf596 = buf588
        del buf588
        buf612 = buf606
        del buf606
        buf628 = buf622
        del buf622
        buf644 = buf638
        del buf638
        triton_poi_fused__softmax_eq_masked_fill_11[grid(64)](buf596,
            buf612, buf628, buf644, primals_7, buf594, buf595, buf610,
            buf611, buf626, buf627, buf642, buf643, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf607 = buf643
        del buf643
        triton_poi_fused_bmm_24[grid(16)](buf589, buf607, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf608 = reinterpret_tensor(buf642, (4, 1, 4), (4, 16, 1), 0)
        del buf642
        triton_poi_fused_bmm_24[grid(16)](buf590, buf608, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf609 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf607, buf608, out=buf609)
        buf623 = reinterpret_tensor(buf608, (4, 4, 1), (4, 1, 16), 0)
        del buf608
        triton_poi_fused_bmm_25[grid(16)](buf589, buf623, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf624 = reinterpret_tensor(buf607, (4, 1, 4), (4, 16, 1), 0)
        del buf607
        triton_poi_fused_bmm_25[grid(16)](buf590, buf624, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf625 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf623, buf624, out=buf625)
        buf639 = reinterpret_tensor(buf624, (4, 4, 1), (4, 1, 16), 0)
        del buf624
        triton_poi_fused_bmm_26[grid(16)](buf589, buf639, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf640 = reinterpret_tensor(buf623, (4, 1, 4), (4, 16, 1), 0)
        del buf623
        triton_poi_fused_bmm_26[grid(16)](buf590, buf640, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf641 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(buf639, buf640, out=buf641)
        buf597 = reinterpret_tensor(buf640, (4, 4, 1), (4, 1, 16), 0)
        del buf640
        buf598 = buf639
        del buf639
        buf613 = buf627
        del buf627
        buf614 = buf626
        del buf626
        buf629 = buf611
        del buf611
        buf630 = buf610
        del buf610
        buf645 = buf595
        del buf595
        buf646 = buf594
        del buf594
        triton_poi_fused__softmax_eq_masked_fill_10[grid(16)](primals_8,
            buf593, buf609, buf625, buf641, buf597, buf598, buf613, buf614,
            buf629, buf630, buf645, buf646, 16, XBLOCK=16, num_warps=1,
            num_stages=1)
        buf599 = buf593
        del buf593
        buf615 = buf609
        del buf609
        buf631 = buf625
        del buf625
        buf647 = buf641
        del buf641
        triton_poi_fused__softmax_eq_masked_fill_11[grid(64)](buf599,
            buf615, buf631, buf647, primals_8, buf597, buf598, buf613,
            buf614, buf629, buf630, buf645, buf646, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf600 = buf646
        del buf646
        triton_poi_fused_bmm_13[grid(16)](buf582, buf600, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf601 = reinterpret_tensor(buf645, (4, 4, 1), (4, 1, 1), 0)
        del buf645
        extern_kernels.bmm(buf596, buf600, out=buf601)
        buf602 = buf600
        del buf600
        triton_poi_fused_bmm_13[grid(16)](buf583, buf602, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf603 = reinterpret_tensor(buf630, (4, 4, 1), (4, 1, 1), 0)
        del buf630
        extern_kernels.bmm(buf599, buf602, out=buf603)
        buf616 = buf602
        del buf602
        triton_poi_fused_bmm_14[grid(16)](buf582, buf616, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf617 = reinterpret_tensor(buf629, (4, 4, 1), (4, 1, 1), 0)
        del buf629
        extern_kernels.bmm(buf612, buf616, out=buf617)
        buf618 = buf616
        del buf616
        triton_poi_fused_bmm_14[grid(16)](buf583, buf618, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf619 = reinterpret_tensor(buf614, (4, 4, 1), (4, 1, 1), 0)
        del buf614
        extern_kernels.bmm(buf615, buf618, out=buf619)
        buf632 = buf618
        del buf618
        triton_poi_fused_bmm_15[grid(16)](buf582, buf632, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf633 = reinterpret_tensor(buf613, (4, 4, 1), (4, 1, 1), 0)
        del buf613
        extern_kernels.bmm(buf628, buf632, out=buf633)
        buf634 = buf632
        del buf632
        triton_poi_fused_bmm_15[grid(16)](buf583, buf634, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf635 = reinterpret_tensor(buf598, (4, 4, 1), (4, 1, 1), 0)
        del buf598
        extern_kernels.bmm(buf631, buf634, out=buf635)
        buf648 = buf634
        del buf634
        triton_poi_fused_bmm_16[grid(16)](buf582, buf648, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf649 = reinterpret_tensor(buf597, (4, 4, 1), (4, 1, 1), 0)
        del buf597
        extern_kernels.bmm(buf644, buf648, out=buf649)
        del buf648
        buf654 = reinterpret_tensor(buf570, (4, 4, 4), (16, 4, 1), 0)
        del buf570
        triton_poi_fused_add_cat_28[grid(64)](buf654, buf601, buf617,
            buf633, buf649, primals_14, 64, XBLOCK=64, num_warps=1,
            num_stages=1)
        del buf601
        del buf617
        del primals_14
        buf651 = reinterpret_tensor(buf649, (4, 4, 1), (4, 1, 16), 0)
        del buf649
        triton_poi_fused_bmm_16[grid(16)](buf583, buf651, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf652 = buf633
        del buf633
        extern_kernels.bmm(buf647, buf651, out=buf652)
        del buf651
        buf656 = reinterpret_tensor(buf571, (4, 4, 4), (16, 4, 1), 0)
        del buf571
        triton_poi_fused_add_cat_28[grid(64)](buf656, buf603, buf619,
            buf635, buf652, primals_16, 64, XBLOCK=64, num_warps=1,
            num_stages=1)
        del buf603
        del buf619
        del buf635
        del buf652
        del primals_16
        buf655 = empty_strided_cuda((16, 4), (4, 1), torch.float32)
        extern_kernels.addmm(primals_26, reinterpret_tensor(buf654, (16, 4),
            (4, 1), 0), reinterpret_tensor(primals_25, (4, 4), (1, 4), 0),
            alpha=1, beta=1, out=buf655)
        del primals_26
        buf657 = empty_strided_cuda((16, 4), (4, 1), torch.float32)
        extern_kernels.addmm(primals_28, reinterpret_tensor(buf656, (16, 4),
            (4, 1), 0), reinterpret_tensor(primals_27, (4, 4), (1, 4), 0),
            alpha=1, beta=1, out=buf657)
        del primals_28
    return (reinterpret_tensor(buf655, (4, 4, 4), (16, 4, 1), 0),
        reinterpret_tensor(buf657, (4, 4, 4), (16, 4, 1), 0), primals_7,
        primals_8, reinterpret_tensor(primals_1, (16, 4), (4, 1), 0),
        reinterpret_tensor(primals_4, (16, 4), (4, 1), 0),
        reinterpret_tensor(buf2, (16, 4), (4, 1), 0), reinterpret_tensor(
        buf4, (16, 4), (4, 1), 0), buf16, buf19, buf32, buf35, buf48, buf51,
        buf64, buf67, reinterpret_tensor(buf75, (16, 8), (8, 1), 0),
        reinterpret_tensor(buf77, (16, 8), (8, 1), 0), buf82, buf83, buf84,
        buf85, reinterpret_tensor(buf86, (16, 4), (4, 1), 0),
        reinterpret_tensor(buf88, (16, 4), (4, 1), 0), reinterpret_tensor(
        buf90, (4, 4, 4), (48, 12, 1), 0), reinterpret_tensor(buf90, (4, 4,
        4), (48, 12, 1), 4), reinterpret_tensor(buf91, (4, 4, 4), (48, 12, 
        1), 0), reinterpret_tensor(buf91, (4, 4, 4), (48, 12, 1), 4),
        buf104, buf107, buf120, buf123, buf136, buf139, buf152, buf155,
        reinterpret_tensor(buf162, (16, 4), (4, 1), 0), reinterpret_tensor(
        buf164, (16, 4), (4, 1), 0), reinterpret_tensor(buf166, (16, 4), (4,
        1), 0), reinterpret_tensor(buf168, (16, 4), (4, 1), 0), buf180,
        buf183, buf196, buf199, buf212, buf215, buf228, buf231,
        reinterpret_tensor(buf239, (16, 8), (8, 1), 0), reinterpret_tensor(
        buf241, (16, 8), (8, 1), 0), buf246, buf247, buf248, buf249,
        reinterpret_tensor(buf250, (16, 4), (4, 1), 0), reinterpret_tensor(
        buf252, (16, 4), (4, 1), 0), reinterpret_tensor(buf254, (4, 4, 4),
        (48, 12, 1), 0), reinterpret_tensor(buf254, (4, 4, 4), (48, 12, 1),
        4), reinterpret_tensor(buf255, (4, 4, 4), (48, 12, 1), 0),
        reinterpret_tensor(buf255, (4, 4, 4), (48, 12, 1), 4), buf268,
        buf271, buf284, buf287, buf300, buf303, buf316, buf319,
        reinterpret_tensor(buf326, (16, 4), (4, 1), 0), reinterpret_tensor(
        buf328, (16, 4), (4, 1), 0), reinterpret_tensor(buf330, (16, 4), (4,
        1), 0), reinterpret_tensor(buf332, (16, 4), (4, 1), 0), buf344,
        buf347, buf360, buf363, buf376, buf379, buf392, buf395,
        reinterpret_tensor(buf403, (16, 8), (8, 1), 0), reinterpret_tensor(
        buf405, (16, 8), (8, 1), 0), buf410, buf411, buf412, buf413,
        reinterpret_tensor(buf414, (16, 4), (4, 1), 0), reinterpret_tensor(
        buf416, (16, 4), (4, 1), 0), reinterpret_tensor(buf418, (4, 4, 4),
        (48, 12, 1), 0), reinterpret_tensor(buf418, (4, 4, 4), (48, 12, 1),
        4), reinterpret_tensor(buf419, (4, 4, 4), (48, 12, 1), 0),
        reinterpret_tensor(buf419, (4, 4, 4), (48, 12, 1), 4), buf432,
        buf435, buf448, buf451, buf464, buf467, buf480, buf483,
        reinterpret_tensor(buf490, (16, 4), (4, 1), 0), reinterpret_tensor(
        buf492, (16, 4), (4, 1), 0), reinterpret_tensor(buf494, (16, 4), (4,
        1), 0), reinterpret_tensor(buf496, (16, 4), (4, 1), 0), buf508,
        buf511, buf524, buf527, buf540, buf543, buf556, buf559,
        reinterpret_tensor(buf567, (16, 8), (8, 1), 0), reinterpret_tensor(
        buf569, (16, 8), (8, 1), 0), buf574, buf575, buf576, buf577,
        reinterpret_tensor(buf578, (16, 4), (4, 1), 0), reinterpret_tensor(
        buf580, (16, 4), (4, 1), 0), reinterpret_tensor(buf582, (4, 4, 4),
        (48, 12, 1), 0), reinterpret_tensor(buf582, (4, 4, 4), (48, 12, 1),
        4), reinterpret_tensor(buf583, (4, 4, 4), (48, 12, 1), 0),
        reinterpret_tensor(buf583, (4, 4, 4), (48, 12, 1), 4), buf596,
        buf599, buf612, buf615, buf628, buf631, buf644, buf647,
        reinterpret_tensor(buf654, (16, 4), (4, 1), 0), reinterpret_tensor(
        buf656, (16, 4), (4, 1), 0), primals_27, primals_25,
        reinterpret_tensor(buf583, (4, 1, 4), (48, 1, 12), 11),
        reinterpret_tensor(buf582, (4, 1, 4), (48, 1, 12), 11),
        reinterpret_tensor(buf589, (4, 1, 4), (16, 1, 4), 3),
        reinterpret_tensor(buf590, (4, 4, 1), (16, 4, 1), 3),
        reinterpret_tensor(buf584, (4, 1, 4), (16, 1, 4), 3),
        reinterpret_tensor(buf585, (4, 4, 1), (16, 4, 1), 3),
        reinterpret_tensor(buf583, (4, 1, 4), (48, 1, 12), 10),
        reinterpret_tensor(buf582, (4, 1, 4), (48, 1, 12), 10),
        reinterpret_tensor(buf589, (4, 1, 4), (16, 1, 4), 2),
        reinterpret_tensor(buf590, (4, 4, 1), (16, 4, 1), 2),
        reinterpret_tensor(buf584, (4, 1, 4), (16, 1, 4), 2),
        reinterpret_tensor(buf585, (4, 4, 1), (16, 4, 1), 2),
        reinterpret_tensor(buf583, (4, 1, 4), (48, 1, 12), 9),
        reinterpret_tensor(buf582, (4, 1, 4), (48, 1, 12), 9),
        reinterpret_tensor(buf589, (4, 1, 4), (16, 1, 4), 1),
        reinterpret_tensor(buf590, (4, 4, 1), (16, 4, 1), 1),
        reinterpret_tensor(buf584, (4, 1, 4), (16, 1, 4), 1),
        reinterpret_tensor(buf585, (4, 4, 1), (16, 4, 1), 1),
        reinterpret_tensor(buf583, (4, 1, 4), (48, 1, 12), 8),
        reinterpret_tensor(buf582, (4, 1, 4), (48, 1, 12), 8),
        reinterpret_tensor(buf589, (4, 1, 4), (16, 1, 4), 0),
        reinterpret_tensor(buf590, (4, 4, 1), (16, 4, 1), 0),
        reinterpret_tensor(buf584, (4, 1, 4), (16, 1, 4), 0),
        reinterpret_tensor(buf585, (4, 4, 1), (16, 4, 1), 0), primals_23,
        buf658, primals_21, buf659, primals_19, primals_17, primals_15,
        primals_13, reinterpret_tensor(buf498, (4, 1, 4), (48, 1, 12), 11),
        reinterpret_tensor(buf499, (4, 1, 4), (48, 1, 12), 11),
        reinterpret_tensor(buf499, (4, 1, 4), (48, 1, 12), 7),
        reinterpret_tensor(buf498, (4, 4, 1), (48, 12, 1), 3),
        reinterpret_tensor(buf498, (4, 1, 4), (48, 1, 12), 7),
        reinterpret_tensor(buf499, (4, 4, 1), (48, 12, 1), 3),
        reinterpret_tensor(buf498, (4, 1, 4), (48, 1, 12), 10),
        reinterpret_tensor(buf499, (4, 1, 4), (48, 1, 12), 10),
        reinterpret_tensor(buf499, (4, 1, 4), (48, 1, 12), 6),
        reinterpret_tensor(buf498, (4, 4, 1), (48, 12, 1), 2),
        reinterpret_tensor(buf498, (4, 1, 4), (48, 1, 12), 6),
        reinterpret_tensor(buf499, (4, 4, 1), (48, 12, 1), 2),
        reinterpret_tensor(buf498, (4, 1, 4), (48, 1, 12), 9),
        reinterpret_tensor(buf499, (4, 1, 4), (48, 1, 12), 9),
        reinterpret_tensor(buf499, (4, 1, 4), (48, 1, 12), 5),
        reinterpret_tensor(buf498, (4, 4, 1), (48, 12, 1), 1),
        reinterpret_tensor(buf498, (4, 1, 4), (48, 1, 12), 5),
        reinterpret_tensor(buf499, (4, 4, 1), (48, 12, 1), 1),
        reinterpret_tensor(buf498, (4, 1, 4), (48, 1, 12), 8),
        reinterpret_tensor(buf499, (4, 1, 4), (48, 1, 12), 8),
        reinterpret_tensor(buf499, (4, 1, 4), (48, 1, 12), 4),
        reinterpret_tensor(buf498, (4, 4, 1), (48, 12, 1), 0),
        reinterpret_tensor(buf498, (4, 1, 4), (48, 1, 12), 4),
        reinterpret_tensor(buf499, (4, 4, 1), (48, 12, 1), 0), primals_11,
        buf660, primals_9, buf661, reinterpret_tensor(buf419, (4, 1, 4), (
        48, 1, 12), 11), reinterpret_tensor(buf418, (4, 1, 4), (48, 1, 12),
        11), reinterpret_tensor(buf425, (4, 1, 4), (16, 1, 4), 3),
        reinterpret_tensor(buf426, (4, 4, 1), (16, 4, 1), 3),
        reinterpret_tensor(buf420, (4, 1, 4), (16, 1, 4), 3),
        reinterpret_tensor(buf421, (4, 4, 1), (16, 4, 1), 3),
        reinterpret_tensor(buf419, (4, 1, 4), (48, 1, 12), 10),
        reinterpret_tensor(buf418, (4, 1, 4), (48, 1, 12), 10),
        reinterpret_tensor(buf425, (4, 1, 4), (16, 1, 4), 2),
        reinterpret_tensor(buf426, (4, 4, 1), (16, 4, 1), 2),
        reinterpret_tensor(buf420, (4, 1, 4), (16, 1, 4), 2),
        reinterpret_tensor(buf421, (4, 4, 1), (16, 4, 1), 2),
        reinterpret_tensor(buf419, (4, 1, 4), (48, 1, 12), 9),
        reinterpret_tensor(buf418, (4, 1, 4), (48, 1, 12), 9),
        reinterpret_tensor(buf425, (4, 1, 4), (16, 1, 4), 1),
        reinterpret_tensor(buf426, (4, 4, 1), (16, 4, 1), 1),
        reinterpret_tensor(buf420, (4, 1, 4), (16, 1, 4), 1),
        reinterpret_tensor(buf421, (4, 4, 1), (16, 4, 1), 1),
        reinterpret_tensor(buf419, (4, 1, 4), (48, 1, 12), 8),
        reinterpret_tensor(buf418, (4, 1, 4), (48, 1, 12), 8),
        reinterpret_tensor(buf425, (4, 1, 4), (16, 1, 4), 0),
        reinterpret_tensor(buf426, (4, 4, 1), (16, 4, 1), 0),
        reinterpret_tensor(buf420, (4, 1, 4), (16, 1, 4), 0),
        reinterpret_tensor(buf421, (4, 4, 1), (16, 4, 1), 0), buf662,
        buf663, reinterpret_tensor(buf334, (4, 1, 4), (48, 1, 12), 11),
        reinterpret_tensor(buf335, (4, 1, 4), (48, 1, 12), 11),
        reinterpret_tensor(buf335, (4, 1, 4), (48, 1, 12), 7),
        reinterpret_tensor(buf334, (4, 4, 1), (48, 12, 1), 3),
        reinterpret_tensor(buf334, (4, 1, 4), (48, 1, 12), 7),
        reinterpret_tensor(buf335, (4, 4, 1), (48, 12, 1), 3),
        reinterpret_tensor(buf334, (4, 1, 4), (48, 1, 12), 10),
        reinterpret_tensor(buf335, (4, 1, 4), (48, 1, 12), 10),
        reinterpret_tensor(buf335, (4, 1, 4), (48, 1, 12), 6),
        reinterpret_tensor(buf334, (4, 4, 1), (48, 12, 1), 2),
        reinterpret_tensor(buf334, (4, 1, 4), (48, 1, 12), 6),
        reinterpret_tensor(buf335, (4, 4, 1), (48, 12, 1), 2),
        reinterpret_tensor(buf334, (4, 1, 4), (48, 1, 12), 9),
        reinterpret_tensor(buf335, (4, 1, 4), (48, 1, 12), 9),
        reinterpret_tensor(buf335, (4, 1, 4), (48, 1, 12), 5),
        reinterpret_tensor(buf334, (4, 4, 1), (48, 12, 1), 1),
        reinterpret_tensor(buf334, (4, 1, 4), (48, 1, 12), 5),
        reinterpret_tensor(buf335, (4, 4, 1), (48, 12, 1), 1),
        reinterpret_tensor(buf334, (4, 1, 4), (48, 1, 12), 8),
        reinterpret_tensor(buf335, (4, 1, 4), (48, 1, 12), 8),
        reinterpret_tensor(buf335, (4, 1, 4), (48, 1, 12), 4),
        reinterpret_tensor(buf334, (4, 4, 1), (48, 12, 1), 0),
        reinterpret_tensor(buf334, (4, 1, 4), (48, 1, 12), 4),
        reinterpret_tensor(buf335, (4, 4, 1), (48, 12, 1), 0), buf664,
        buf665, reinterpret_tensor(buf255, (4, 1, 4), (48, 1, 12), 11),
        reinterpret_tensor(buf254, (4, 1, 4), (48, 1, 12), 11),
        reinterpret_tensor(buf261, (4, 1, 4), (16, 1, 4), 3),
        reinterpret_tensor(buf262, (4, 4, 1), (16, 4, 1), 3),
        reinterpret_tensor(buf256, (4, 1, 4), (16, 1, 4), 3),
        reinterpret_tensor(buf257, (4, 4, 1), (16, 4, 1), 3),
        reinterpret_tensor(buf255, (4, 1, 4), (48, 1, 12), 10),
        reinterpret_tensor(buf254, (4, 1, 4), (48, 1, 12), 10),
        reinterpret_tensor(buf261, (4, 1, 4), (16, 1, 4), 2),
        reinterpret_tensor(buf262, (4, 4, 1), (16, 4, 1), 2),
        reinterpret_tensor(buf256, (4, 1, 4), (16, 1, 4), 2),
        reinterpret_tensor(buf257, (4, 4, 1), (16, 4, 1), 2),
        reinterpret_tensor(buf255, (4, 1, 4), (48, 1, 12), 9),
        reinterpret_tensor(buf254, (4, 1, 4), (48, 1, 12), 9),
        reinterpret_tensor(buf261, (4, 1, 4), (16, 1, 4), 1),
        reinterpret_tensor(buf262, (4, 4, 1), (16, 4, 1), 1),
        reinterpret_tensor(buf256, (4, 1, 4), (16, 1, 4), 1),
        reinterpret_tensor(buf257, (4, 4, 1), (16, 4, 1), 1),
        reinterpret_tensor(buf255, (4, 1, 4), (48, 1, 12), 8),
        reinterpret_tensor(buf254, (4, 1, 4), (48, 1, 12), 8),
        reinterpret_tensor(buf261, (4, 1, 4), (16, 1, 4), 0),
        reinterpret_tensor(buf262, (4, 4, 1), (16, 4, 1), 0),
        reinterpret_tensor(buf256, (4, 1, 4), (16, 1, 4), 0),
        reinterpret_tensor(buf257, (4, 4, 1), (16, 4, 1), 0), buf666,
        buf667, reinterpret_tensor(buf170, (4, 1, 4), (48, 1, 12), 11),
        reinterpret_tensor(buf171, (4, 1, 4), (48, 1, 12), 11),
        reinterpret_tensor(buf171, (4, 1, 4), (48, 1, 12), 7),
        reinterpret_tensor(buf170, (4, 4, 1), (48, 12, 1), 3),
        reinterpret_tensor(buf170, (4, 1, 4), (48, 1, 12), 7),
        reinterpret_tensor(buf171, (4, 4, 1), (48, 12, 1), 3),
        reinterpret_tensor(buf170, (4, 1, 4), (48, 1, 12), 10),
        reinterpret_tensor(buf171, (4, 1, 4), (48, 1, 12), 10),
        reinterpret_tensor(buf171, (4, 1, 4), (48, 1, 12), 6),
        reinterpret_tensor(buf170, (4, 4, 1), (48, 12, 1), 2),
        reinterpret_tensor(buf170, (4, 1, 4), (48, 1, 12), 6),
        reinterpret_tensor(buf171, (4, 4, 1), (48, 12, 1), 2),
        reinterpret_tensor(buf170, (4, 1, 4), (48, 1, 12), 9),
        reinterpret_tensor(buf171, (4, 1, 4), (48, 1, 12), 9),
        reinterpret_tensor(buf171, (4, 1, 4), (48, 1, 12), 5),
        reinterpret_tensor(buf170, (4, 4, 1), (48, 12, 1), 1),
        reinterpret_tensor(buf170, (4, 1, 4), (48, 1, 12), 5),
        reinterpret_tensor(buf171, (4, 4, 1), (48, 12, 1), 1),
        reinterpret_tensor(buf170, (4, 1, 4), (48, 1, 12), 8),
        reinterpret_tensor(buf171, (4, 1, 4), (48, 1, 12), 8),
        reinterpret_tensor(buf171, (4, 1, 4), (48, 1, 12), 4),
        reinterpret_tensor(buf170, (4, 4, 1), (48, 12, 1), 0),
        reinterpret_tensor(buf170, (4, 1, 4), (48, 1, 12), 4),
        reinterpret_tensor(buf171, (4, 4, 1), (48, 12, 1), 0), buf668,
        buf669, reinterpret_tensor(buf91, (4, 1, 4), (48, 1, 12), 11),
        reinterpret_tensor(buf90, (4, 1, 4), (48, 1, 12), 11),
        reinterpret_tensor(buf97, (4, 1, 4), (16, 1, 4), 3),
        reinterpret_tensor(buf98, (4, 4, 1), (16, 4, 1), 3),
        reinterpret_tensor(buf92, (4, 1, 4), (16, 1, 4), 3),
        reinterpret_tensor(buf93, (4, 4, 1), (16, 4, 1), 3),
        reinterpret_tensor(buf91, (4, 1, 4), (48, 1, 12), 10),
        reinterpret_tensor(buf90, (4, 1, 4), (48, 1, 12), 10),
        reinterpret_tensor(buf97, (4, 1, 4), (16, 1, 4), 2),
        reinterpret_tensor(buf98, (4, 4, 1), (16, 4, 1), 2),
        reinterpret_tensor(buf92, (4, 1, 4), (16, 1, 4), 2),
        reinterpret_tensor(buf93, (4, 4, 1), (16, 4, 1), 2),
        reinterpret_tensor(buf91, (4, 1, 4), (48, 1, 12), 9),
        reinterpret_tensor(buf90, (4, 1, 4), (48, 1, 12), 9),
        reinterpret_tensor(buf97, (4, 1, 4), (16, 1, 4), 1),
        reinterpret_tensor(buf98, (4, 4, 1), (16, 4, 1), 1),
        reinterpret_tensor(buf92, (4, 1, 4), (16, 1, 4), 1),
        reinterpret_tensor(buf93, (4, 4, 1), (16, 4, 1), 1),
        reinterpret_tensor(buf91, (4, 1, 4), (48, 1, 12), 8),
        reinterpret_tensor(buf90, (4, 1, 4), (48, 1, 12), 8),
        reinterpret_tensor(buf97, (4, 1, 4), (16, 1, 4), 0),
        reinterpret_tensor(buf98, (4, 4, 1), (16, 4, 1), 0),
        reinterpret_tensor(buf92, (4, 1, 4), (16, 1, 4), 0),
        reinterpret_tensor(buf93, (4, 4, 1), (16, 4, 1), 0), buf670, buf671,
        reinterpret_tensor(buf6, (4, 1, 4), (48, 1, 12), 11),
        reinterpret_tensor(buf7, (4, 1, 4), (48, 1, 12), 11),
        reinterpret_tensor(buf7, (4, 1, 4), (48, 1, 12), 7),
        reinterpret_tensor(buf6, (4, 4, 1), (48, 12, 1), 3),
        reinterpret_tensor(buf6, (4, 1, 4), (48, 1, 12), 7),
        reinterpret_tensor(buf7, (4, 4, 1), (48, 12, 1), 3),
        reinterpret_tensor(buf6, (4, 1, 4), (48, 1, 12), 10),
        reinterpret_tensor(buf7, (4, 1, 4), (48, 1, 12), 10),
        reinterpret_tensor(buf7, (4, 1, 4), (48, 1, 12), 6),
        reinterpret_tensor(buf6, (4, 4, 1), (48, 12, 1), 2),
        reinterpret_tensor(buf6, (4, 1, 4), (48, 1, 12), 6),
        reinterpret_tensor(buf7, (4, 4, 1), (48, 12, 1), 2),
        reinterpret_tensor(buf6, (4, 1, 4), (48, 1, 12), 9),
        reinterpret_tensor(buf7, (4, 1, 4), (48, 1, 12), 9),
        reinterpret_tensor(buf7, (4, 1, 4), (48, 1, 12), 5),
        reinterpret_tensor(buf6, (4, 4, 1), (48, 12, 1), 1),
        reinterpret_tensor(buf6, (4, 1, 4), (48, 1, 12), 5),
        reinterpret_tensor(buf7, (4, 4, 1), (48, 12, 1), 1),
        reinterpret_tensor(buf6, (4, 1, 4), (48, 1, 12), 8),
        reinterpret_tensor(buf7, (4, 1, 4), (48, 1, 12), 8),
        reinterpret_tensor(buf7, (4, 1, 4), (48, 1, 12), 4),
        reinterpret_tensor(buf6, (4, 4, 1), (48, 12, 1), 0),
        reinterpret_tensor(buf6, (4, 1, 4), (48, 1, 12), 4),
        reinterpret_tensor(buf7, (4, 4, 1), (48, 12, 1), 0), buf672, buf673)


class DyIntraModalityUpdate(nn.Module):
    """
    Dynamic Intra-modality Attention Flow
    """

    def __init__(self, v_size, q_size, output_size, num_head, drop=0.0):
        super(DyIntraModalityUpdate, self).__init__()
        self.v_size = v_size
        self.q_size = q_size
        self.output_size = output_size
        self.num_head = num_head
        self.v4q_gate_lin = nn.Linear(v_size, output_size)
        self.q4v_gate_lin = nn.Linear(q_size, output_size)
        self.v_lin = nn.Linear(v_size, output_size * 3)
        self.q_lin = nn.Linear(q_size, output_size * 3)
        self.v_output = nn.Linear(output_size, output_size)
        self.q_output = nn.Linear(output_size, output_size)
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()
        self.drop = nn.Dropout(drop)

    def forward(self, v, q, v_mask, q_mask):
        """
        v: visual feature      [batch, num_obj, feat_size]
        q: question            [batch, max_len, feat_size]
        v_mask                 [batch, num_obj]
        q_mask                 [batch, max_len]
        """
        batch_size, num_obj = v_mask.shape
        _, max_len = q_mask.shape
        v_mean = (v * v_mask.unsqueeze(2)).sum(1) / v_mask.sum(1).unsqueeze(1)
        q_mean = (q * q_mask.unsqueeze(2)).sum(1) / q_mask.sum(1).unsqueeze(1)
        v4q_gate = self.sigmoid(self.v4q_gate_lin(self.drop(self.relu(v_mean)))
            ).unsqueeze(1)
        q4v_gate = self.sigmoid(self.q4v_gate_lin(self.drop(self.relu(q_mean)))
            ).unsqueeze(1)
        v_trans = self.v_lin(self.drop(self.relu(v)))
        q_trans = self.q_lin(self.drop(self.relu(q)))
        v_trans = v_trans * v_mask.unsqueeze(2)
        q_trans = q_trans * q_mask.unsqueeze(2)
        v_k, v_q, v_v = torch.split(v_trans, v_trans.size(2) // 3, dim=2)
        q_k, q_q, q_v = torch.split(q_trans, q_trans.size(2) // 3, dim=2)
        new_vq = (1 + q4v_gate) * v_q
        new_vk = (1 + q4v_gate) * v_k
        new_qq = (1 + v4q_gate) * q_q
        new_qk = (1 + v4q_gate) * q_k
        vk_set = torch.split(new_vk, new_vk.size(2) // self.num_head, dim=2)
        vq_set = torch.split(new_vq, new_vq.size(2) // self.num_head, dim=2)
        vv_set = torch.split(v_v, v_v.size(2) // self.num_head, dim=2)
        qk_set = torch.split(new_qk, new_qk.size(2) // self.num_head, dim=2)
        qq_set = torch.split(new_qq, new_qq.size(2) // self.num_head, dim=2)
        qv_set = torch.split(q_v, q_v.size(2) // self.num_head, dim=2)
        for i in range(self.num_head):
            vk_slice, vq_slice, vv_slice = vk_set[i], vq_set[i], vv_set[i]
            qk_slice, qq_slice, qv_slice = qk_set[i], qq_set[i], qv_set[i]
            v2v = (vq_slice @ vk_slice.transpose(1, 2)).masked_fill(v_mask.
                unsqueeze(1).expand([batch_size, num_obj, num_obj]) == 0, -
                1000000000.0) / (self.output_size // self.num_head) ** 0.5
            q2q = (qq_slice @ qk_slice.transpose(1, 2)).masked_fill(q_mask.
                unsqueeze(1).expand([batch_size, max_len, max_len]) == 0, -
                1000000000.0) / (self.output_size // self.num_head) ** 0.5
            dyIntraMAF_v2v = F.softmax(v2v, dim=2)
            dyIntraMAF_q2q = F.softmax(q2q, dim=2)
            v_update = dyIntraMAF_v2v @ vv_slice if i == 0 else torch.cat((
                v_update, dyIntraMAF_v2v @ vv_slice), dim=2)
            q_update = dyIntraMAF_q2q @ qv_slice if i == 0 else torch.cat((
                q_update, dyIntraMAF_q2q @ qv_slice), dim=2)
        updated_v = self.v_output(self.drop(v + v_update))
        updated_q = self.q_output(self.drop(q + q_update))
        return updated_v, updated_q


class InterModalityUpdate(nn.Module):
    """
    Inter-modality Attention Flow
    """

    def __init__(self, v_size, q_size, output_size, num_head, drop=0.0):
        super(InterModalityUpdate, self).__init__()
        self.v_size = v_size
        self.q_size = q_size
        self.output_size = output_size
        self.num_head = num_head
        self.v_lin = nn.Linear(v_size, output_size * 3)
        self.q_lin = nn.Linear(q_size, output_size * 3)
        self.v_output = nn.Linear(output_size + v_size, output_size)
        self.q_output = nn.Linear(output_size + q_size, output_size)
        self.relu = nn.ReLU()
        self.drop = nn.Dropout(drop)

    def forward(self, v, q, v_mask, q_mask):
        """
        v: visual feature      [batch, num_obj, feat_size]
        q: question            [batch, max_len, feat_size]
        v_mask                 [batch, num_obj]
        q_mask                 [batch, max_len]
        """
        batch_size, num_obj = v_mask.shape
        _, max_len = q_mask.shape
        v_trans = self.v_lin(self.drop(self.relu(v)))
        q_trans = self.q_lin(self.drop(self.relu(q)))
        v_trans = v_trans * v_mask.unsqueeze(2)
        q_trans = q_trans * q_mask.unsqueeze(2)
        v_k, v_q, v_v = torch.split(v_trans, v_trans.size(2) // 3, dim=2)
        q_k, q_q, q_v = torch.split(q_trans, q_trans.size(2) // 3, dim=2)
        vk_set = torch.split(v_k, v_k.size(2) // self.num_head, dim=2)
        vq_set = torch.split(v_q, v_q.size(2) // self.num_head, dim=2)
        vv_set = torch.split(v_v, v_v.size(2) // self.num_head, dim=2)
        qk_set = torch.split(q_k, q_k.size(2) // self.num_head, dim=2)
        qq_set = torch.split(q_q, q_q.size(2) // self.num_head, dim=2)
        qv_set = torch.split(q_v, q_v.size(2) // self.num_head, dim=2)
        for i in range(self.num_head):
            vk_slice, vq_slice, vv_slice = vk_set[i], vq_set[i], vv_set[i]
            qk_slice, qq_slice, qv_slice = qk_set[i], qq_set[i], qv_set[i]
            q2v = (vq_slice @ qk_slice.transpose(1, 2)).masked_fill(q_mask.
                unsqueeze(1).expand([batch_size, num_obj, max_len]) == 0, -
                1000000000.0) / (self.output_size // self.num_head) ** 0.5
            v2q = (qq_slice @ vk_slice.transpose(1, 2)).masked_fill(v_mask.
                unsqueeze(1).expand([batch_size, max_len, num_obj]) == 0, -
                1000000000.0) / (self.output_size // self.num_head) ** 0.5
            interMAF_q2v = F.softmax(q2v, dim=2)
            interMAF_v2q = F.softmax(v2q, dim=2)
            v_update = interMAF_q2v @ qv_slice if i == 0 else torch.cat((
                v_update, interMAF_q2v @ qv_slice), dim=2)
            q_update = interMAF_v2q @ vv_slice if i == 0 else torch.cat((
                q_update, interMAF_v2q @ vv_slice), dim=2)
        cat_v = torch.cat((v, v_update), dim=2)
        cat_q = torch.cat((q, q_update), dim=2)
        updated_v = self.v_output(self.drop(cat_v))
        updated_q = self.q_output(self.drop(cat_q))
        return updated_v, updated_q


class SingleBlockNew(nn.Module):
    """
    Single Block Inter-/Intra-modality stack multiple times
    """

    def __init__(self, num_block, v_size, q_size, output_size,
        num_inter_head, num_intra_head, drop=0.0):
        super(SingleBlockNew, self).__init__()
        self.v_size = v_size
        self.q_size = q_size
        self.output_size = output_size
        self.num_inter_head = num_inter_head
        self.num_intra_head = num_intra_head
        self.num_block = num_block
        self.v_lin = nn.Linear(v_size, output_size)
        self.q_lin = nn.Linear(q_size, output_size)
        self.interBlock = InterModalityUpdate(output_size, output_size,
            output_size, num_inter_head, drop)
        self.intraBlock = DyIntraModalityUpdate(output_size, output_size,
            output_size, num_intra_head, drop)
        self.drop = nn.Dropout(drop)

    def forward(self, input_0, input_1, input_2, input_3):
        primals_2 = self.v_lin.weight
        primals_3 = self.v_lin.bias
        primals_5 = self.q_lin.weight
        primals_6 = self.q_lin.bias
        primals_9 = self.interBlock.v_lin.weight
        primals_10 = self.interBlock.v_lin.bias
        primals_11 = self.interBlock.q_lin.weight
        primals_12 = self.interBlock.q_lin.bias
        primals_13 = self.interBlock.v_output.weight
        primals_14 = self.interBlock.v_output.bias
        primals_15 = self.interBlock.q_output.weight
        primals_16 = self.interBlock.q_output.bias
        primals_7 = self.intraBlock.v4q_gate_lin.weight
        primals_18 = self.intraBlock.v4q_gate_lin.bias
        primals_8 = self.intraBlock.q4v_gate_lin.weight
        primals_20 = self.intraBlock.q4v_gate_lin.bias
        primals_21 = self.intraBlock.v_lin.weight
        primals_22 = self.intraBlock.v_lin.bias
        primals_23 = self.intraBlock.q_lin.weight
        primals_24 = self.intraBlock.q_lin.bias
        primals_17 = self.intraBlock.v_output.weight
        primals_26 = self.intraBlock.v_output.bias
        primals_19 = self.intraBlock.q_output.weight
        primals_28 = self.intraBlock.q_output.bias
        primals_1 = input_0
        primals_4 = input_1
        primals_25 = input_2
        primals_27 = input_3
        output = call([primals_1, primals_2, primals_3, primals_4,
            primals_5, primals_6, primals_7, primals_8, primals_9,
            primals_10, primals_11, primals_12, primals_13, primals_14,
            primals_15, primals_16, primals_17, primals_18, primals_19,
            primals_20, primals_21, primals_22, primals_23, primals_24,
            primals_25, primals_26, primals_27, primals_28])
        return output[0], output[1]


Measured latency_ms: 0.015216