storylm0_tp_proj.kina

let proj_q0 = fn(
  x: tensor<F32, [1, 4, 4], @host>,
  wq: tensor<F32, [4, 4], @host>
) -> tensor<F32, [1, 1, 4, 2], @host> {
  let wq0 = wq.slice(axis=1, start=0, stop=2) in
  let wq0_b = wq0.broadcast([1, 4, 2]) in
  let q_flat0 = x @ wq0_b in
  let q0 = q_flat0.reshape([1, 4, 1, 2]).transpose([0, 2, 1, 3]) in
  q0
} in
let tp_proj = fn(
  x: tensor<F32, [1, 4, 4], @host>,
  wq: tensor<F32, [4, 4], @host>,
  wk: tensor<F32, [4, 4], @host>,
  wv: tensor<F32, [4, 4], @host>,
  attention_scale: tensor<F32, [1, 1, 1, 1], @host>,
  causal_mask: tensor<F32, [1, 1, 4, 4], @host>
) -> tensor<F32, [1, 2, 4, 2], @host> {
  let wq0 = wq.slice(axis=1, start=0, stop=2) in
  let wq1 = wq.slice(axis=1, start=2, stop=4) in
  let wk0 = wk.slice(axis=1, start=0, stop=2) in
  let wk1 = wk.slice(axis=1, start=2, stop=4) in
  let wv0 = wv.slice(axis=1, start=0, stop=2) in
  let wv1 = wv.slice(axis=1, start=2, stop=4) in
  let wq0_b = wq0.broadcast([1, 4, 2]) in
  let wq1_b = wq1.broadcast([1, 4, 2]) in
  let wk0_b = wk0.broadcast([1, 4, 2]) in
  let wk1_b = wk1.broadcast([1, 4, 2]) in
  let wv0_b = wv0.broadcast([1, 4, 2]) in
  let wv1_b = wv1.broadcast([1, 4, 2]) in
  let q_flat0 = x @ wq0_b in
  let q_flat1 = x @ wq1_b in
  let k_flat0 = x @ wk0_b in
  let k_flat1 = x @ wk1_b in
  let v_flat0 = x @ wv0_b in
  let v_flat1 = x @ wv1_b in
  let q0 = q_flat0.reshape([1, 4, 1, 2]).transpose([0, 2, 1, 3]) in
  let q1 = q_flat1.reshape([1, 4, 1, 2]).transpose([0, 2, 1, 3]) in
  let k0 = k_flat0.reshape([1, 4, 1, 2]).transpose([0, 2, 3, 1]) in
  let k1 = k_flat1.reshape([1, 4, 1, 2]).transpose([0, 2, 3, 1]) in
  let v0 = v_flat0.reshape([1, 4, 1, 2]).transpose([0, 2, 1, 3]) in
  let v1 = v_flat1.reshape([1, 4, 1, 2]).transpose([0, 2, 1, 3]) in
  let scale_b = attention_scale.broadcast([1, 1, 4, 4]) in
  let mask_b = causal_mask.broadcast([1, 1, 4, 4]) in
  let scores0 = q0 @ k0 in
  let scores1 = q1 @ k1 in
  let probs0 = ((scores0 * scale_b) + mask_b).softmax(axis=3) in
  let probs1 = ((scores1 * scale_b) + mask_b).softmax(axis=3) in
  let ctx0 = probs0 @ v0 in
  let ctx1 = probs1 @ v1 in
  let context = ctx0.concat(axis=1, ctx1) in
  context
} in
let storylm0_tp_proj = fn(
  token_ids: tensor<I64, [1, 4], @host>,
  token_embedding: tensor<F32, [8, 4], @host>,
  causal_mask: tensor<F32, [1, 1, 4, 4], @host>,
  attention_scale: tensor<F32, [1, 1, 1, 1], @host>,
  wq: tensor<F32, [4, 4], @host>,
  wk: tensor<F32, [4, 4], @host>,
  wv: tensor<F32, [4, 4], @host>,
  wo: tensor<F32, [4, 4], @host>,
  lm_head: tensor<F32, [4, 8], @host>
) -> tensor<F32, [1, 4, 8], @host> {
  let x = token_embedding.gather(token_ids) in
  let context = tp_proj(x, wq, wk, wv, attention_scale, causal_mask) in
  let wo_broadcast = wo.broadcast([1, 4, 4]) in
  let lm_head_broadcast = lm_head.broadcast([1, 4, 8]) in
  let merged = context.transpose([0, 2, 1, 3]).reshape([1, 4, 4]) in
  let projected = merged @ wo_broadcast in
  let residual = x + projected in
  let logits = residual @ lm_head_broadcast in
  logits
} in
let token_ids = tensor<I64, [1, 4], @host> in
let token_embedding = tensor<F32, [8, 4], @host> in
let causal_mask = tensor<F32, [1, 1, 4, 4], @host> in
let attention_scale = tensor<F32, [1, 1, 1, 1], @host> in
let wq = tensor<F32, [4, 4], @host> in
let wk = tensor<F32, [4, 4], @host> in
let wv = tensor<F32, [4, 4], @host> in
let wo = tensor<F32, [4, 4], @host> in
let lm_head = tensor<F32, [4, 8], @host> in
storylm0_tp_proj(token_ids, token_embedding, causal_mask, attention_scale, wq, wk, wv, wo, lm_head)

module {
  func.func @fn_val_0(%1: tensor<1x32x64xf32>, %2: tensor<1x1x64xf32>) -> tensor<1x32x64xf32> attributes {llvm.emit_c_interface} {
      %3 = tensor.empty() : tensor<1x32x64xf32>
      %4 = linalg.mul ins(%1, %1 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%3 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
      %7 = arith.constant 0.000000e+00 : f32
      %5 = tensor.empty() : tensor<1x32x1xf32>
      %6 = linalg.fill ins(%7 : f32) outs(%5 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
      %8 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
        iterator_types = ["parallel", "parallel", "reduction"]
      } ins(%4 : tensor<1x32x64xf32>) outs(%6 : tensor<1x32x1xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        %9 = arith.addf %a0, %a1 : f32
        linalg.yield %9 : f32
      } -> tensor<1x32x1xf32>
      %10 = arith.constant dense<0.000000e+00> : tensor<1x1x1xf32>
      %11 = tensor.empty() : tensor<1x32x1xf32>
      %12 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
        iterator_types = ["parallel", "parallel", "parallel"]
      } ins(%10 : tensor<1x1x1xf32>) outs(%11 : tensor<1x32x1xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x32x1xf32>
      %13 = tensor.empty() : tensor<1x32x1xf32>
      %14 = linalg.mul ins(%8, %12 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%13 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
      %15 = arith.constant dense<0.000000e+00> : tensor<1x1x1xf32>
      %16 = tensor.empty() : tensor<1x32x1xf32>
      %17 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
        iterator_types = ["parallel", "parallel", "parallel"]
      } ins(%15 : tensor<1x1x1xf32>) outs(%16 : tensor<1x32x1xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x32x1xf32>
      %18 = tensor.empty() : tensor<1x32x1xf32>
      %19 = linalg.add ins(%14, %17 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%18 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
      %23 = tensor.empty() : tensor<1x32x1xf32>
      %24 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
        iterator_types = ["parallel", "parallel", "parallel"]
      } ins(%19 : tensor<1x32x1xf32>) outs(%23 : tensor<1x32x1xf32>) {
      ^bb0(%x: f32, %_: f32):
        %20 = math.sqrt %x : f32
        %21 = arith.constant 1.000000e+00 : f32
        %22 = arith.divf %21, %20 : f32
        linalg.yield %22 : f32
      } -> tensor<1x32x1xf32>
      %25 = tensor.empty() : tensor<1x32x64xf32>
      %26 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
        iterator_types = ["parallel", "parallel", "parallel"]
      } ins(%24 : tensor<1x32x1xf32>) outs(%25 : tensor<1x32x64xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x32x64xf32>
      %27 = tensor.empty() : tensor<1x32x64xf32>
      %28 = linalg.mul ins(%1, %26 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%27 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
      %29 = tensor.empty() : tensor<1x32x64xf32>
      %30 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
        iterator_types = ["parallel", "parallel", "parallel"]
      } ins(%2 : tensor<1x1x64xf32>) outs(%29 : tensor<1x32x64xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x32x64xf32>
      %31 = tensor.empty() : tensor<1x32x64xf32>
      %32 = linalg.mul ins(%28, %30 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%31 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
    return %32 : tensor<1x32x64xf32>
  }
  func.func @fn_val_33(%34: tensor<1x32x64xf32>, %35: tensor<1x64x128xf32>, %36: tensor<1x128x64xf32>) -> tensor<1x32x64xf32> attributes {llvm.emit_c_interface} {
      %37 = arith.constant dense<0.000000e+00> : tensor<1x32x128xf32>
      %38 = linalg.batch_matmul ins(%34, %35 : tensor<1x32x64xf32>, tensor<1x64x128xf32>) outs(%37 : tensor<1x32x128xf32>) -> tensor<1x32x128xf32>
      %44 = tensor.empty() : tensor<1x32x128xf32>
      %45 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
        iterator_types = ["parallel", "parallel", "parallel"]
      } ins(%38 : tensor<1x32x128xf32>) outs(%44 : tensor<1x32x128xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        %39 = arith.constant 1.000000e+00 : f32
        %40 = arith.negf %a0 : f32
        %41 = math.exp %40 : f32
        %42 = arith.addf %39, %41 : f32
        %43 = arith.divf %39, %42 : f32
        linalg.yield %43 : f32
      } -> tensor<1x32x128xf32>
      %46 = tensor.empty() : tensor<1x32x128xf32>
      %47 = linalg.mul ins(%38, %45 : tensor<1x32x128xf32>, tensor<1x32x128xf32>) outs(%46 : tensor<1x32x128xf32>) -> tensor<1x32x128xf32>
      %48 = arith.constant dense<0.000000e+00> : tensor<1x32x64xf32>
      %49 = linalg.batch_matmul ins(%47, %36 : tensor<1x32x128xf32>, tensor<1x128x64xf32>) outs(%48 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
    return %49 : tensor<1x32x64xf32>
  }
  func.func @fn_val_50(%51: tensor<1x32x64xf32>, %52: tensor<1x64x64xf32>, %53: tensor<1x64x64xf32>, %54: tensor<1x64x64xf32>, %55: tensor<1x64x64xf32>) -> tensor<1x32x64xf32> attributes {llvm.emit_c_interface} {
      %56 = arith.constant dense<0.000000e+00> : tensor<1x32x64xf32>
      %57 = linalg.batch_matmul ins(%51, %52 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%56 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
      %58 = arith.constant dense<0.000000e+00> : tensor<1x32x64xf32>
      %59 = linalg.batch_matmul ins(%51, %53 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%58 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
      %60 = arith.constant dense<0.000000e+00> : tensor<1x32x64xf32>
      %61 = linalg.batch_matmul ins(%51, %54 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%60 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
      %62 = tensor.empty() : tensor<1x64x32xf32>
      %63 = linalg.transpose ins(%59 : tensor<1x32x64xf32>) outs(%62 : tensor<1x64x32xf32>) permutation = [0, 2, 1]
      %64 = arith.constant dense<0.000000e+00> : tensor<1x32x32xf32>
      %65 = linalg.batch_matmul ins(%57, %63 : tensor<1x32x64xf32>, tensor<1x64x32xf32>) outs(%64 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
      %66 = arith.constant dense<0.000000e+00> : tensor<1x1x1xf32>
      %67 = tensor.empty() : tensor<1x32x32xf32>
      %68 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
        iterator_types = ["parallel", "parallel", "parallel"]
      } ins(%66 : tensor<1x1x1xf32>) outs(%67 : tensor<1x32x32xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x32x32xf32>
      %69 = tensor.empty() : tensor<1x32x32xf32>
      %70 = linalg.mul ins(%65, %68 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%69 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
      %73 = arith.constant 0xFF800000 : f32
      %71 = tensor.empty() : tensor<1x32x1xf32>
      %72 = linalg.fill ins(%73 : f32) outs(%71 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
      %74 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
        iterator_types = ["parallel", "parallel", "reduction"]
      } ins(%70 : tensor<1x32x32xf32>) outs(%72 : tensor<1x32x1xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        %75 = arith.maximumf %a0, %a1 : f32
        linalg.yield %75 : f32
      } -> tensor<1x32x1xf32>
      %76 = tensor.empty() : tensor<1x32x32xf32>
      %77 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
        iterator_types = ["parallel", "parallel", "parallel"]
      } ins(%74 : tensor<1x32x1xf32>) outs(%76 : tensor<1x32x32xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x32x32xf32>
      %78 = tensor.empty() : tensor<1x32x32xf32>
      %79 = linalg.sub ins(%70, %77 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%78 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
      %80 = tensor.empty() : tensor<1x32x32xf32>
      %81 = linalg.exp ins(%79 : tensor<1x32x32xf32>) outs(%80 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
      %84 = arith.constant 0xFF800000 : f32
      %82 = tensor.empty() : tensor<1x32x1xf32>
      %83 = linalg.fill ins(%84 : f32) outs(%82 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
      %85 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
        iterator_types = ["parallel", "parallel", "reduction"]
      } ins(%70 : tensor<1x32x32xf32>) outs(%83 : tensor<1x32x1xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        %86 = arith.maximumf %a0, %a1 : f32
        linalg.yield %86 : f32
      } -> tensor<1x32x1xf32>
      %87 = tensor.empty() : tensor<1x32x32xf32>
      %88 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
        iterator_types = ["parallel", "parallel", "parallel"]
      } ins(%85 : tensor<1x32x1xf32>) outs(%87 : tensor<1x32x32xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x32x32xf32>
      %89 = tensor.empty() : tensor<1x32x32xf32>
      %90 = linalg.sub ins(%70, %88 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%89 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
      %91 = tensor.empty() : tensor<1x32x32xf32>
      %92 = linalg.exp ins(%90 : tensor<1x32x32xf32>) outs(%91 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
      %95 = arith.constant 0.000000e+00 : f32
      %93 = tensor.empty() : tensor<1x32x1xf32>
      %94 = linalg.fill ins(%95 : f32) outs(%93 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
      %96 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
        iterator_types = ["parallel", "parallel", "reduction"]
      } ins(%92 : tensor<1x32x32xf32>) outs(%94 : tensor<1x32x1xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        %97 = arith.addf %a0, %a1 : f32
        linalg.yield %97 : f32
      } -> tensor<1x32x1xf32>
      %98 = tensor.empty() : tensor<1x32x32xf32>
      %99 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
        iterator_types = ["parallel", "parallel", "parallel"]
      } ins(%96 : tensor<1x32x1xf32>) outs(%98 : tensor<1x32x32xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x32x32xf32>
      %100 = tensor.empty() : tensor<1x32x32xf32>
      %101 = linalg.div ins(%81, %99 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%100 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
      %102 = arith.constant dense<0.000000e+00> : tensor<1x32x64xf32>
      %103 = linalg.batch_matmul ins(%101, %61 : tensor<1x32x32xf32>, tensor<1x32x64xf32>) outs(%102 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
      %104 = arith.constant dense<0.000000e+00> : tensor<1x32x64xf32>
      %105 = linalg.batch_matmul ins(%103, %55 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%104 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
    return %105 : tensor<1x32x64xf32>
  }
  func.func @fn_val_106(%107: tensor<1x32x64xf32>, %108: tensor<1x1x64xf32>, %109: tensor<1x64x64xf32>, %110: tensor<1x64x64xf32>, %111: tensor<1x64x64xf32>, %112: tensor<1x64x64xf32>, %113: tensor<1x1x64xf32>, %114: tensor<1x64x128xf32>, %115: tensor<1x128x64xf32>) -> tensor<1x32x64xf32> attributes {llvm.emit_c_interface} {
      %116 = func.call @fn_val_0(%107, %108) : (tensor<1x32x64xf32>, tensor<1x1x64xf32>) -> tensor<1x32x64xf32>
      %117 = func.call @fn_val_50(%116, %109, %110, %111, %112) : (tensor<1x32x64xf32>, tensor<1x64x64xf32>, tensor<1x64x64xf32>, tensor<1x64x64xf32>, tensor<1x64x64xf32>) -> tensor<1x32x64xf32>
      %118 = tensor.empty() : tensor<1x32x64xf32>
      %119 = linalg.add ins(%107, %117 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%118 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
      %120 = func.call @fn_val_0(%119, %113) : (tensor<1x32x64xf32>, tensor<1x1x64xf32>) -> tensor<1x32x64xf32>
      %121 = func.call @fn_val_33(%120, %114, %115) : (tensor<1x32x64xf32>, tensor<1x64x128xf32>, tensor<1x128x64xf32>) -> tensor<1x32x64xf32>
      %122 = tensor.empty() : tensor<1x32x64xf32>
      %123 = linalg.add ins(%119, %121 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%122 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
    return %123 : tensor<1x32x64xf32>
  }
  func.func @fn_val_124(%125: tensor<1x4x4xf32>, %126: tensor<4x4xf32>) -> tensor<1x1x4x2xf32> attributes {llvm.emit_c_interface} {
      %127 = tensor.extract_slice %126[0, 0][4, 2][1, 1] : tensor<4x4xf32> to tensor<4x2xf32>
      %128 = tensor.empty() : tensor<1x4x2xf32>
      %129 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
        iterator_types = ["parallel", "parallel", "parallel"]
      } ins(%127 : tensor<4x2xf32>) outs(%128 : tensor<1x4x2xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x4x2xf32>
      %130 = arith.constant dense<0.000000e+00> : tensor<1x4x2xf32>
      %131 = linalg.batch_matmul ins(%125, %129 : tensor<1x4x4xf32>, tensor<1x4x2xf32>) outs(%130 : tensor<1x4x2xf32>) -> tensor<1x4x2xf32>
      %132 = arith.constant dense<[1, 4, 1, 2]> : tensor<4xi32>
      %133 = tensor.reshape %131(%132) : (tensor<1x4x2xf32>, tensor<4xi32>) -> tensor<1x4x1x2xf32>
      %134 = tensor.empty() : tensor<1x1x4x2xf32>
      %135 = linalg.transpose ins(%133 : tensor<1x4x1x2xf32>) outs(%134 : tensor<1x1x4x2xf32>) permutation = [0, 2, 1, 3]
    return %135 : tensor<1x1x4x2xf32>
  }
  func.func @fn_val_136(%137: tensor<1x4x4xf32>, %138: tensor<4x4xf32>, %139: tensor<4x4xf32>, %140: tensor<4x4xf32>, %141: tensor<1x1x1x1xf32>, %142: tensor<1x1x4x4xf32>) -> tensor<1x2x4x2xf32> attributes {llvm.emit_c_interface} {
      %143 = tensor.extract_slice %138[0, 0][4, 2][1, 1] : tensor<4x4xf32> to tensor<4x2xf32>
      %144 = tensor.extract_slice %138[0, 2][4, 2][1, 1] : tensor<4x4xf32> to tensor<4x2xf32>
      %145 = tensor.extract_slice %139[0, 0][4, 2][1, 1] : tensor<4x4xf32> to tensor<4x2xf32>
      %146 = tensor.extract_slice %139[0, 2][4, 2][1, 1] : tensor<4x4xf32> to tensor<4x2xf32>
      %147 = tensor.extract_slice %140[0, 0][4, 2][1, 1] : tensor<4x4xf32> to tensor<4x2xf32>
      %148 = tensor.extract_slice %140[0, 2][4, 2][1, 1] : tensor<4x4xf32> to tensor<4x2xf32>
      %149 = tensor.empty() : tensor<1x4x2xf32>
      %150 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
        iterator_types = ["parallel", "parallel", "parallel"]
      } ins(%143 : tensor<4x2xf32>) outs(%149 : tensor<1x4x2xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x4x2xf32>
      %151 = tensor.empty() : tensor<1x4x2xf32>
      %152 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
        iterator_types = ["parallel", "parallel", "parallel"]
      } ins(%144 : tensor<4x2xf32>) outs(%151 : tensor<1x4x2xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x4x2xf32>
      %153 = tensor.empty() : tensor<1x4x2xf32>
      %154 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
        iterator_types = ["parallel", "parallel", "parallel"]
      } ins(%145 : tensor<4x2xf32>) outs(%153 : tensor<1x4x2xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x4x2xf32>
      %155 = tensor.empty() : tensor<1x4x2xf32>
      %156 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
        iterator_types = ["parallel", "parallel", "parallel"]
      } ins(%146 : tensor<4x2xf32>) outs(%155 : tensor<1x4x2xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x4x2xf32>
      %157 = tensor.empty() : tensor<1x4x2xf32>
      %158 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
        iterator_types = ["parallel", "parallel", "parallel"]
      } ins(%147 : tensor<4x2xf32>) outs(%157 : tensor<1x4x2xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x4x2xf32>
      %159 = tensor.empty() : tensor<1x4x2xf32>
      %160 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
        iterator_types = ["parallel", "parallel", "parallel"]
      } ins(%148 : tensor<4x2xf32>) outs(%159 : tensor<1x4x2xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x4x2xf32>
      %161 = arith.constant dense<0.000000e+00> : tensor<1x4x2xf32>
      %162 = linalg.batch_matmul ins(%137, %150 : tensor<1x4x4xf32>, tensor<1x4x2xf32>) outs(%161 : tensor<1x4x2xf32>) -> tensor<1x4x2xf32>
      %163 = arith.constant dense<0.000000e+00> : tensor<1x4x2xf32>
      %164 = linalg.batch_matmul ins(%137, %152 : tensor<1x4x4xf32>, tensor<1x4x2xf32>) outs(%163 : tensor<1x4x2xf32>) -> tensor<1x4x2xf32>
      %165 = arith.constant dense<0.000000e+00> : tensor<1x4x2xf32>
      %166 = linalg.batch_matmul ins(%137, %154 : tensor<1x4x4xf32>, tensor<1x4x2xf32>) outs(%165 : tensor<1x4x2xf32>) -> tensor<1x4x2xf32>
      %167 = arith.constant dense<0.000000e+00> : tensor<1x4x2xf32>
      %168 = linalg.batch_matmul ins(%137, %156 : tensor<1x4x4xf32>, tensor<1x4x2xf32>) outs(%167 : tensor<1x4x2xf32>) -> tensor<1x4x2xf32>
      %169 = arith.constant dense<0.000000e+00> : tensor<1x4x2xf32>
      %170 = linalg.batch_matmul ins(%137, %158 : tensor<1x4x4xf32>, tensor<1x4x2xf32>) outs(%169 : tensor<1x4x2xf32>) -> tensor<1x4x2xf32>
      %171 = arith.constant dense<0.000000e+00> : tensor<1x4x2xf32>
      %172 = linalg.batch_matmul ins(%137, %160 : tensor<1x4x4xf32>, tensor<1x4x2xf32>) outs(%171 : tensor<1x4x2xf32>) -> tensor<1x4x2xf32>
      %173 = arith.constant dense<[1, 4, 1, 2]> : tensor<4xi32>
      %174 = tensor.reshape %162(%173) : (tensor<1x4x2xf32>, tensor<4xi32>) -> tensor<1x4x1x2xf32>
      %175 = tensor.empty() : tensor<1x1x4x2xf32>
      %176 = linalg.transpose ins(%174 : tensor<1x4x1x2xf32>) outs(%175 : tensor<1x1x4x2xf32>) permutation = [0, 2, 1, 3]
      %177 = arith.constant dense<[1, 4, 1, 2]> : tensor<4xi32>
      %178 = tensor.reshape %164(%177) : (tensor<1x4x2xf32>, tensor<4xi32>) -> tensor<1x4x1x2xf32>
      %179 = tensor.empty() : tensor<1x1x4x2xf32>
      %180 = linalg.transpose ins(%178 : tensor<1x4x1x2xf32>) outs(%179 : tensor<1x1x4x2xf32>) permutation = [0, 2, 1, 3]
      %181 = arith.constant dense<[1, 4, 1, 2]> : tensor<4xi32>
      %182 = tensor.reshape %166(%181) : (tensor<1x4x2xf32>, tensor<4xi32>) -> tensor<1x4x1x2xf32>
      %183 = tensor.empty() : tensor<1x1x2x4xf32>
      %184 = linalg.transpose ins(%182 : tensor<1x4x1x2xf32>) outs(%183 : tensor<1x1x2x4xf32>) permutation = [0, 2, 3, 1]
      %185 = arith.constant dense<[1, 4, 1, 2]> : tensor<4xi32>
      %186 = tensor.reshape %168(%185) : (tensor<1x4x2xf32>, tensor<4xi32>) -> tensor<1x4x1x2xf32>
      %187 = tensor.empty() : tensor<1x1x2x4xf32>
      %188 = linalg.transpose ins(%186 : tensor<1x4x1x2xf32>) outs(%187 : tensor<1x1x2x4xf32>) permutation = [0, 2, 3, 1]
      %189 = arith.constant dense<[1, 4, 1, 2]> : tensor<4xi32>
      %190 = tensor.reshape %170(%189) : (tensor<1x4x2xf32>, tensor<4xi32>) -> tensor<1x4x1x2xf32>
      %191 = tensor.empty() : tensor<1x1x4x2xf32>
      %192 = linalg.transpose ins(%190 : tensor<1x4x1x2xf32>) outs(%191 : tensor<1x1x4x2xf32>) permutation = [0, 2, 1, 3]
      %193 = arith.constant dense<[1, 4, 1, 2]> : tensor<4xi32>
      %194 = tensor.reshape %172(%193) : (tensor<1x4x2xf32>, tensor<4xi32>) -> tensor<1x4x1x2xf32>
      %195 = tensor.empty() : tensor<1x1x4x2xf32>
      %196 = linalg.transpose ins(%194 : tensor<1x4x1x2xf32>) outs(%195 : tensor<1x1x4x2xf32>) permutation = [0, 2, 1, 3]
      %197 = tensor.empty() : tensor<1x1x4x4xf32>
      %198 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, 0, 0, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
        iterator_types = ["parallel", "parallel", "parallel", "parallel"]
      } ins(%141 : tensor<1x1x1x1xf32>) outs(%197 : tensor<1x1x4x4xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x1x4x4xf32>
      %199 = tensor.empty() : tensor<1x1x4x4xf32>
      %200 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, 0, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
        iterator_types = ["parallel", "parallel", "parallel", "parallel"]
      } ins(%142 : tensor<1x1x4x4xf32>) outs(%199 : tensor<1x1x4x4xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x1x4x4xf32>
      %201 = tensor.collapse_shape %176 [[0, 1], [2], [3]] : tensor<1x1x4x2xf32> into tensor<1x4x2xf32>
      %202 = tensor.collapse_shape %184 [[0, 1], [2], [3]] : tensor<1x1x2x4xf32> into tensor<1x2x4xf32>
      %203 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
      %204 = linalg.batch_matmul ins(%201, %202 : tensor<1x4x2xf32>, tensor<1x2x4xf32>) outs(%203 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
      %205 = tensor.expand_shape %204 [[0, 1], [2], [3]] : tensor<1x4x4xf32> into tensor<1x1x4x4xf32>
      %206 = tensor.collapse_shape %180 [[0, 1], [2], [3]] : tensor<1x1x4x2xf32> into tensor<1x4x2xf32>
      %207 = tensor.collapse_shape %188 [[0, 1], [2], [3]] : tensor<1x1x2x4xf32> into tensor<1x2x4xf32>
      %208 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
      %209 = linalg.batch_matmul ins(%206, %207 : tensor<1x4x2xf32>, tensor<1x2x4xf32>) outs(%208 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
      %210 = tensor.expand_shape %209 [[0, 1], [2], [3]] : tensor<1x4x4xf32> into tensor<1x1x4x4xf32>
      %211 = tensor.empty() : tensor<1x1x4x4xf32>
      %212 = linalg.mul ins(%205, %198 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%211 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %213 = tensor.empty() : tensor<1x1x4x4xf32>
      %214 = linalg.add ins(%212, %200 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%213 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %215 = tensor.empty() : tensor<1x1x4x4xf32>
      %216 = linalg.mul ins(%205, %198 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%215 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %217 = tensor.empty() : tensor<1x1x4x4xf32>
      %218 = linalg.add ins(%216, %200 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%217 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %221 = arith.constant 0xFF800000 : f32
      %219 = tensor.empty() : tensor<1x1x4x1xf32>
      %220 = linalg.fill ins(%221 : f32) outs(%219 : tensor<1x1x4x1xf32>) -> tensor<1x1x4x1xf32>
      %222 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, 0)>],
        iterator_types = ["parallel", "parallel", "parallel", "reduction"]
      } ins(%218 : tensor<1x1x4x4xf32>) outs(%220 : tensor<1x1x4x1xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        %223 = arith.maximumf %a0, %a1 : f32
        linalg.yield %223 : f32
      } -> tensor<1x1x4x1xf32>
      %224 = tensor.empty() : tensor<1x1x4x4xf32>
      %225 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, 0, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
        iterator_types = ["parallel", "parallel", "parallel", "parallel"]
      } ins(%222 : tensor<1x1x4x1xf32>) outs(%224 : tensor<1x1x4x4xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x1x4x4xf32>
      %226 = tensor.empty() : tensor<1x1x4x4xf32>
      %227 = linalg.sub ins(%214, %225 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%226 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %228 = tensor.empty() : tensor<1x1x4x4xf32>
      %229 = linalg.exp ins(%227 : tensor<1x1x4x4xf32>) outs(%228 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %230 = tensor.empty() : tensor<1x1x4x4xf32>
      %231 = linalg.mul ins(%205, %198 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%230 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %232 = tensor.empty() : tensor<1x1x4x4xf32>
      %233 = linalg.add ins(%231, %200 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%232 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %234 = tensor.empty() : tensor<1x1x4x4xf32>
      %235 = linalg.mul ins(%205, %198 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%234 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %236 = tensor.empty() : tensor<1x1x4x4xf32>
      %237 = linalg.add ins(%235, %200 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%236 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %240 = arith.constant 0xFF800000 : f32
      %238 = tensor.empty() : tensor<1x1x4x1xf32>
      %239 = linalg.fill ins(%240 : f32) outs(%238 : tensor<1x1x4x1xf32>) -> tensor<1x1x4x1xf32>
      %241 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, 0)>],
        iterator_types = ["parallel", "parallel", "parallel", "reduction"]
      } ins(%237 : tensor<1x1x4x4xf32>) outs(%239 : tensor<1x1x4x1xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        %242 = arith.maximumf %a0, %a1 : f32
        linalg.yield %242 : f32
      } -> tensor<1x1x4x1xf32>
      %243 = tensor.empty() : tensor<1x1x4x4xf32>
      %244 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, 0, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
        iterator_types = ["parallel", "parallel", "parallel", "parallel"]
      } ins(%241 : tensor<1x1x4x1xf32>) outs(%243 : tensor<1x1x4x4xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x1x4x4xf32>
      %245 = tensor.empty() : tensor<1x1x4x4xf32>
      %246 = linalg.sub ins(%233, %244 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%245 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %247 = tensor.empty() : tensor<1x1x4x4xf32>
      %248 = linalg.exp ins(%246 : tensor<1x1x4x4xf32>) outs(%247 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %251 = arith.constant 0.000000e+00 : f32
      %249 = tensor.empty() : tensor<1x1x4x1xf32>
      %250 = linalg.fill ins(%251 : f32) outs(%249 : tensor<1x1x4x1xf32>) -> tensor<1x1x4x1xf32>
      %252 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, 0)>],
        iterator_types = ["parallel", "parallel", "parallel", "reduction"]
      } ins(%248 : tensor<1x1x4x4xf32>) outs(%250 : tensor<1x1x4x1xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        %253 = arith.addf %a0, %a1 : f32
        linalg.yield %253 : f32
      } -> tensor<1x1x4x1xf32>
      %254 = tensor.empty() : tensor<1x1x4x4xf32>
      %255 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, 0, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
        iterator_types = ["parallel", "parallel", "parallel", "parallel"]
      } ins(%252 : tensor<1x1x4x1xf32>) outs(%254 : tensor<1x1x4x4xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x1x4x4xf32>
      %256 = tensor.empty() : tensor<1x1x4x4xf32>
      %257 = linalg.div ins(%229, %255 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%256 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %258 = tensor.empty() : tensor<1x1x4x4xf32>
      %259 = linalg.mul ins(%210, %198 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%258 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %260 = tensor.empty() : tensor<1x1x4x4xf32>
      %261 = linalg.add ins(%259, %200 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%260 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %262 = tensor.empty() : tensor<1x1x4x4xf32>
      %263 = linalg.mul ins(%210, %198 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%262 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %264 = tensor.empty() : tensor<1x1x4x4xf32>
      %265 = linalg.add ins(%263, %200 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%264 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %268 = arith.constant 0xFF800000 : f32
      %266 = tensor.empty() : tensor<1x1x4x1xf32>
      %267 = linalg.fill ins(%268 : f32) outs(%266 : tensor<1x1x4x1xf32>) -> tensor<1x1x4x1xf32>
      %269 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, 0)>],
        iterator_types = ["parallel", "parallel", "parallel", "reduction"]
      } ins(%265 : tensor<1x1x4x4xf32>) outs(%267 : tensor<1x1x4x1xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        %270 = arith.maximumf %a0, %a1 : f32
        linalg.yield %270 : f32
      } -> tensor<1x1x4x1xf32>
      %271 = tensor.empty() : tensor<1x1x4x4xf32>
      %272 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, 0, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
        iterator_types = ["parallel", "parallel", "parallel", "parallel"]
      } ins(%269 : tensor<1x1x4x1xf32>) outs(%271 : tensor<1x1x4x4xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x1x4x4xf32>
      %273 = tensor.empty() : tensor<1x1x4x4xf32>
      %274 = linalg.sub ins(%261, %272 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%273 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %275 = tensor.empty() : tensor<1x1x4x4xf32>
      %276 = linalg.exp ins(%274 : tensor<1x1x4x4xf32>) outs(%275 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %277 = tensor.empty() : tensor<1x1x4x4xf32>
      %278 = linalg.mul ins(%210, %198 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%277 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %279 = tensor.empty() : tensor<1x1x4x4xf32>
      %280 = linalg.add ins(%278, %200 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%279 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %281 = tensor.empty() : tensor<1x1x4x4xf32>
      %282 = linalg.mul ins(%210, %198 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%281 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %283 = tensor.empty() : tensor<1x1x4x4xf32>
      %284 = linalg.add ins(%282, %200 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%283 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %287 = arith.constant 0xFF800000 : f32
      %285 = tensor.empty() : tensor<1x1x4x1xf32>
      %286 = linalg.fill ins(%287 : f32) outs(%285 : tensor<1x1x4x1xf32>) -> tensor<1x1x4x1xf32>
      %288 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, 0)>],
        iterator_types = ["parallel", "parallel", "parallel", "reduction"]
      } ins(%284 : tensor<1x1x4x4xf32>) outs(%286 : tensor<1x1x4x1xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        %289 = arith.maximumf %a0, %a1 : f32
        linalg.yield %289 : f32
      } -> tensor<1x1x4x1xf32>
      %290 = tensor.empty() : tensor<1x1x4x4xf32>
      %291 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, 0, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
        iterator_types = ["parallel", "parallel", "parallel", "parallel"]
      } ins(%288 : tensor<1x1x4x1xf32>) outs(%290 : tensor<1x1x4x4xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x1x4x4xf32>
      %292 = tensor.empty() : tensor<1x1x4x4xf32>
      %293 = linalg.sub ins(%280, %291 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%292 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %294 = tensor.empty() : tensor<1x1x4x4xf32>
      %295 = linalg.exp ins(%293 : tensor<1x1x4x4xf32>) outs(%294 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %298 = arith.constant 0.000000e+00 : f32
      %296 = tensor.empty() : tensor<1x1x4x1xf32>
      %297 = linalg.fill ins(%298 : f32) outs(%296 : tensor<1x1x4x1xf32>) -> tensor<1x1x4x1xf32>
      %299 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, 0)>],
        iterator_types = ["parallel", "parallel", "parallel", "reduction"]
      } ins(%295 : tensor<1x1x4x4xf32>) outs(%297 : tensor<1x1x4x1xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        %300 = arith.addf %a0, %a1 : f32
        linalg.yield %300 : f32
      } -> tensor<1x1x4x1xf32>
      %301 = tensor.empty() : tensor<1x1x4x4xf32>
      %302 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, 0, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
        iterator_types = ["parallel", "parallel", "parallel", "parallel"]
      } ins(%299 : tensor<1x1x4x1xf32>) outs(%301 : tensor<1x1x4x4xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x1x4x4xf32>
      %303 = tensor.empty() : tensor<1x1x4x4xf32>
      %304 = linalg.div ins(%276, %302 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%303 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
      %305 = tensor.collapse_shape %257 [[0, 1], [2], [3]] : tensor<1x1x4x4xf32> into tensor<1x4x4xf32>
      %306 = tensor.collapse_shape %192 [[0, 1], [2], [3]] : tensor<1x1x4x2xf32> into tensor<1x4x2xf32>
      %307 = arith.constant dense<0.000000e+00> : tensor<1x4x2xf32>
      %308 = linalg.batch_matmul ins(%305, %306 : tensor<1x4x4xf32>, tensor<1x4x2xf32>) outs(%307 : tensor<1x4x2xf32>) -> tensor<1x4x2xf32>
      %309 = tensor.expand_shape %308 [[0, 1], [2], [3]] : tensor<1x4x2xf32> into tensor<1x1x4x2xf32>
      %310 = tensor.collapse_shape %304 [[0, 1], [2], [3]] : tensor<1x1x4x4xf32> into tensor<1x4x4xf32>
      %311 = tensor.collapse_shape %196 [[0, 1], [2], [3]] : tensor<1x1x4x2xf32> into tensor<1x4x2xf32>
      %312 = arith.constant dense<0.000000e+00> : tensor<1x4x2xf32>
      %313 = linalg.batch_matmul ins(%310, %311 : tensor<1x4x4xf32>, tensor<1x4x2xf32>) outs(%312 : tensor<1x4x2xf32>) -> tensor<1x4x2xf32>
      %314 = tensor.expand_shape %313 [[0, 1], [2], [3]] : tensor<1x4x2xf32> into tensor<1x1x4x2xf32>
      %315 = tensor.empty() : tensor<1x2x4x2xf32>
      %316 = tensor.insert_slice %309 into %315[0, 0, 0, 0] [1, 1, 4, 2] [1, 1, 1, 1] : tensor<1x1x4x2xf32> into tensor<1x2x4x2xf32>
      %317 = tensor.insert_slice %314 into %316[0, 1, 0, 0] [1, 1, 4, 2] [1, 1, 1, 1] : tensor<1x1x4x2xf32> into tensor<1x2x4x2xf32>
    return %317 : tensor<1x2x4x2xf32>
  }
  func.func @fn_val_318(%319: tensor<1x4xi64>, %320: tensor<8x4xf32>, %321: tensor<1x1x4x4xf32>, %322: tensor<1x1x1x1xf32>, %323: tensor<4x4xf32>, %324: tensor<4x4xf32>, %325: tensor<4x4xf32>, %326: tensor<4x4xf32>, %327: tensor<4x8xf32>) -> tensor<1x4x8xf32> attributes {llvm.emit_c_interface} {
      %328 = arith.constant 0 : i64
      %329 = arith.constant 8 : i64
      %330 = tensor.empty() : tensor<1x4xi64>
      %331 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
        iterator_types = ["parallel", "parallel"]
      } ins(%319 : tensor<1x4xi64>) outs(%330 : tensor<1x4xi64>) {
      ^bb0(%candidate: i64, %_: i64):
        %332 = arith.cmpi sge, %candidate, %328 : i64
        %333 = arith.cmpi slt, %candidate, %329 : i64
        %334 = arith.andi %332, %333 : i1
        cf.assert %334, "gather index out of bounds"
        linalg.yield %candidate : i64
      } -> tensor<1x4xi64>
      %335 = tensor.empty() : tensor<1x4x4xf32>
      %336 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
        iterator_types = ["parallel", "parallel", "parallel"]
      } ins(%331 : tensor<1x4xi64>) outs(%335 : tensor<1x4x4xf32>) {
      ^bb0(%idx: i64, %_: f32):
        %337 = arith.index_cast %idx : i64 to index
        %338 = linalg.index 2 : index
        %339 = tensor.extract %320[%337, %338] : tensor<8x4xf32>
        linalg.yield %339 : f32
      } -> tensor<1x4x4xf32>
      %340 = func.call @fn_val_136(%336, %323, %324, %325, %322, %321) : (tensor<1x4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<1x1x1x1xf32>, tensor<1x1x4x4xf32>) -> tensor<1x2x4x2xf32>
      %341 = tensor.empty() : tensor<1x4x4xf32>
      %342 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
        iterator_types = ["parallel", "parallel", "parallel"]
      } ins(%326 : tensor<4x4xf32>) outs(%341 : tensor<1x4x4xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x4x4xf32>
      %343 = tensor.empty() : tensor<1x4x8xf32>
      %344 = linalg.generic {
        indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
        iterator_types = ["parallel", "parallel", "parallel"]
      } ins(%327 : tensor<4x8xf32>) outs(%343 : tensor<1x4x8xf32>) {
      ^bb0(%a0: f32, %a1: f32):
        linalg.yield %a0 : f32
      } -> tensor<1x4x8xf32>
      %345 = tensor.empty() : tensor<1x4x2x2xf32>
      %346 = linalg.transpose ins(%340 : tensor<1x2x4x2xf32>) outs(%345 : tensor<1x4x2x2xf32>) permutation = [0, 2, 1, 3]
      %347 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
      %348 = tensor.reshape %346(%347) : (tensor<1x4x2x2xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
      %349 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
      %350 = linalg.batch_matmul ins(%348, %342 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%349 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
      %351 = tensor.empty() : tensor<1x4x4xf32>
      %352 = linalg.add ins(%336, %350 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%351 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
      %353 = arith.constant dense<0.000000e+00> : tensor<1x4x8xf32>
      %354 = linalg.batch_matmul ins(%352, %344 : tensor<1x4x4xf32>, tensor<1x4x8xf32>) outs(%353 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
    return %354 : tensor<1x4x8xf32>
  }
  func.func @main() -> tensor<1x4x8xf32> {
    %355 = arith.constant dense<0> : tensor<1x4xi64>
    %356 = arith.constant dense<0.000000e+00> : tensor<8x4xf32>
    %357 = arith.constant dense<0.000000e+00> : tensor<1x1x4x4xf32>
    %358 = arith.constant dense<0.000000e+00> : tensor<1x1x1x1xf32>
    %359 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
    %360 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
    %361 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
    %362 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
    %363 = arith.constant dense<0.000000e+00> : tensor<4x8xf32>
    %364 = func.call @fn_val_318(%355, %356, %357, %358, %359, %360, %361, %362, %363) : (tensor<1x4xi64>, tensor<8x4xf32>, tensor<1x1x4x4xf32>, tensor<1x1x1x1xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x8xf32>) -> tensor<1x4x8xf32>
    return %364 : tensor<1x4x8xf32>
  }
}