Example
storylm0_nll_step.kina
let storylm0_nll = fn(
token_ids: tensor<I64, [1, 4], @host>,
targets: 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], @host> {
let x = token_embedding.gather(token_ids) in
// Unit batch dim is a volume-preserving reshape. Rank-changing broadcast
// still has no VJP in the restricted pass.
let wq_broadcast = wq.reshape([1, 4, 4]) in
let wk_broadcast = wk.reshape([1, 4, 4]) in
let wv_broadcast = wv.reshape([1, 4, 4]) in
let wo_broadcast = wo.reshape([1, 4, 4]) in
let lm_head_broadcast = lm_head.reshape([1, 4, 8]) in
let q_flat = x @ wq_broadcast in
let k_flat = x @ wk_broadcast in
let v_flat = x @ wv_broadcast in
let q_unsplit = q_flat.reshape([1, 4, 2, 2]) in
let q = q_unsplit.transpose([0, 2, 1, 3]) in
let k_unsplit = k_flat.reshape([1, 4, 2, 2]) in
let k_t = k_unsplit.transpose([0, 2, 3, 1]) in
let v_unsplit = v_flat.reshape([1, 4, 2, 2]) in
let v = v_unsplit.transpose([0, 2, 1, 3]) in
let scores = q @ k_t in
let scale_broadcast = attention_scale.broadcast([1, 2, 4, 4]) in
let mask_broadcast = causal_mask.broadcast([1, 2, 4, 4]) in
let scaled_scores = scores * scale_broadcast in
let masked_scores = scaled_scores + mask_broadcast in
let probs = masked_scores.softmax(axis=3) in
let context = probs @ v in
let context_t = context.transpose([0, 2, 1, 3]) in
let merged = context_t.reshape([1, 4, 4]) in
let projected = merged @ wo_broadcast in
let residual = x + projected in
let logits = residual @ lm_head_broadcast in
let log_probs = logits.log_softmax(axis=2) in
let picked = log_probs.take_along_axis(targets, axis=2) in
let nll = -picked in
nll
} in
let step = fn(
token_ids: tensor<I64, [1, 4], @host>,
targets: 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>,
lr: tensor<F32, [1, 1], @host>
) -> (
tensor<F32, [1, 4], @host>,
tensor<F32, [8, 4], @host>,
tensor<F32, [4, 4], @host>,
tensor<F32, [4, 4], @host>,
tensor<F32, [4, 4], @host>,
tensor<F32, [4, 4], @host>,
tensor<F32, [4, 8], @host>
) {
let (nll, g_emb, g_mask, g_scale, g_wq, g_wk, g_wv, g_wo, g_lm) =
handle storylm0_nll(
token_ids, targets, token_embedding, causal_mask, attention_scale,
wq, wk, wv, wo, lm_head
) with reverseMode in
let lr_emb = lr.broadcast([8, 4]) in
let emb_delta = lr_emb * g_emb in
let emb_next = token_embedding - emb_delta in
let lr_w = lr.broadcast([4, 4]) in
let wq_delta = lr_w * g_wq in
let wq_next = wq - wq_delta in
let wk_delta = lr_w * g_wk in
let wk_next = wk - wk_delta in
let wv_delta = lr_w * g_wv in
let wv_next = wv - wv_delta in
let wo_delta = lr_w * g_wo in
let wo_next = wo - wo_delta in
let lr_lm = lr.broadcast([4, 8]) in
let lm_delta = lr_lm * g_lm in
let lm_next = lm_head - lm_delta in
(nll, emb_next, wq_next, wk_next, wv_next, wo_next, lm_next)
} in
let token_ids = tensor<I64, [1, 4], @host> in
let targets = 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
let lr = tensor<F32, [1, 1], @host> in
step(token_ids, targets, token_embedding, causal_mask, attention_scale, wq, wk, wv, wo, lm_head, lr)
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<1x4xi64>, %126: tensor<1x4xi64>, %127: tensor<8x4xf32>, %128: tensor<1x1x4x4xf32>, %129: tensor<1x1x1x1xf32>, %130: tensor<4x4xf32>, %131: tensor<4x4xf32>, %132: tensor<4x4xf32>, %133: tensor<4x4xf32>, %134: tensor<4x8xf32>) -> tensor<1x4xf32> attributes {llvm.emit_c_interface} {
%135 = arith.constant 0 : i64
%136 = arith.constant 8 : i64
%137 = tensor.empty() : tensor<1x4xi64>
%138 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%125 : tensor<1x4xi64>) outs(%137 : tensor<1x4xi64>) {
^bb0(%candidate: i64, %_: i64):
%139 = arith.cmpi sge, %candidate, %135 : i64
%140 = arith.cmpi slt, %candidate, %136 : i64
%141 = arith.andi %139, %140 : i1
cf.assert %141, "gather index out of bounds"
linalg.yield %candidate : i64
} -> tensor<1x4xi64>
%142 = tensor.empty() : tensor<1x4x4xf32>
%143 = 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(%138 : tensor<1x4xi64>) outs(%142 : tensor<1x4x4xf32>) {
^bb0(%idx: i64, %_: f32):
%144 = arith.index_cast %idx : i64 to index
%145 = linalg.index 2 : index
%146 = tensor.extract %127[%144, %145] : tensor<8x4xf32>
linalg.yield %146 : f32
} -> tensor<1x4x4xf32>
%147 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%148 = tensor.reshape %130(%147) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%149 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%150 = tensor.reshape %131(%149) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%151 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%152 = tensor.reshape %132(%151) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%153 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%154 = tensor.reshape %133(%153) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%155 = arith.constant dense<[1, 4, 8]> : tensor<3xi32>
%156 = tensor.reshape %134(%155) : (tensor<4x8xf32>, tensor<3xi32>) -> tensor<1x4x8xf32>
%157 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%158 = linalg.batch_matmul ins(%143, %148 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%157 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%159 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%160 = linalg.batch_matmul ins(%143, %150 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%159 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%161 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%162 = linalg.batch_matmul ins(%143, %152 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%161 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%163 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%164 = tensor.reshape %158(%163) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%165 = tensor.empty() : tensor<1x2x4x2xf32>
%166 = linalg.transpose ins(%164 : tensor<1x4x2x2xf32>) outs(%165 : tensor<1x2x4x2xf32>) permutation = [0, 2, 1, 3]
%167 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%168 = tensor.reshape %160(%167) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%169 = tensor.empty() : tensor<1x2x2x4xf32>
%170 = linalg.transpose ins(%168 : tensor<1x4x2x2xf32>) outs(%169 : tensor<1x2x2x4xf32>) permutation = [0, 2, 3, 1]
%171 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%172 = tensor.reshape %162(%171) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%173 = tensor.empty() : tensor<1x2x4x2xf32>
%174 = linalg.transpose ins(%172 : tensor<1x4x2x2xf32>) outs(%173 : tensor<1x2x4x2xf32>) permutation = [0, 2, 1, 3]
%175 = tensor.collapse_shape %166 [[0, 1], [2], [3]] : tensor<1x2x4x2xf32> into tensor<2x4x2xf32>
%176 = tensor.collapse_shape %170 [[0, 1], [2], [3]] : tensor<1x2x2x4xf32> into tensor<2x2x4xf32>
%177 = arith.constant dense<0.000000e+00> : tensor<2x4x4xf32>
%178 = linalg.batch_matmul ins(%175, %176 : tensor<2x4x2xf32>, tensor<2x2x4xf32>) outs(%177 : tensor<2x4x4xf32>) -> tensor<2x4x4xf32>
%179 = tensor.expand_shape %178 [[0, 1], [2], [3]] : tensor<2x4x4xf32> into tensor<1x2x4x4xf32>
%180 = tensor.empty() : tensor<1x2x4x4xf32>
%181 = 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(%129 : tensor<1x1x1x1xf32>) outs(%180 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%182 = tensor.empty() : tensor<1x2x4x4xf32>
%183 = 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(%128 : tensor<1x1x4x4xf32>) outs(%182 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%184 = tensor.empty() : tensor<1x2x4x4xf32>
%185 = linalg.mul ins(%179, %181 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%184 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%186 = tensor.empty() : tensor<1x2x4x4xf32>
%187 = linalg.add ins(%185, %183 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%186 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%190 = arith.constant 0xFF800000 : f32
%188 = tensor.empty() : tensor<1x2x4x1xf32>
%189 = linalg.fill ins(%190 : f32) outs(%188 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%191 = 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(%187 : tensor<1x2x4x4xf32>) outs(%189 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%192 = arith.maximumf %a0, %a1 : f32
linalg.yield %192 : f32
} -> tensor<1x2x4x1xf32>
%193 = tensor.empty() : tensor<1x2x4x4xf32>
%194 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, d1, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%191 : tensor<1x2x4x1xf32>) outs(%193 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%195 = tensor.empty() : tensor<1x2x4x4xf32>
%196 = linalg.sub ins(%187, %194 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%195 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%197 = tensor.empty() : tensor<1x2x4x4xf32>
%198 = linalg.exp ins(%196 : tensor<1x2x4x4xf32>) outs(%197 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%201 = arith.constant 0xFF800000 : f32
%199 = tensor.empty() : tensor<1x2x4x1xf32>
%200 = linalg.fill ins(%201 : f32) outs(%199 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%202 = 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(%187 : tensor<1x2x4x4xf32>) outs(%200 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%203 = arith.maximumf %a0, %a1 : f32
linalg.yield %203 : f32
} -> tensor<1x2x4x1xf32>
%204 = tensor.empty() : tensor<1x2x4x4xf32>
%205 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, d1, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%202 : tensor<1x2x4x1xf32>) outs(%204 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%206 = tensor.empty() : tensor<1x2x4x4xf32>
%207 = linalg.sub ins(%187, %205 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%206 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%208 = tensor.empty() : tensor<1x2x4x4xf32>
%209 = linalg.exp ins(%207 : tensor<1x2x4x4xf32>) outs(%208 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%212 = arith.constant 0.000000e+00 : f32
%210 = tensor.empty() : tensor<1x2x4x1xf32>
%211 = linalg.fill ins(%212 : f32) outs(%210 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%213 = 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(%209 : tensor<1x2x4x4xf32>) outs(%211 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%214 = arith.addf %a0, %a1 : f32
linalg.yield %214 : f32
} -> tensor<1x2x4x1xf32>
%215 = tensor.empty() : tensor<1x2x4x4xf32>
%216 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, d1, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%213 : tensor<1x2x4x1xf32>) outs(%215 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%217 = tensor.empty() : tensor<1x2x4x4xf32>
%218 = linalg.div ins(%198, %216 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%217 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%219 = tensor.collapse_shape %218 [[0, 1], [2], [3]] : tensor<1x2x4x4xf32> into tensor<2x4x4xf32>
%220 = tensor.collapse_shape %174 [[0, 1], [2], [3]] : tensor<1x2x4x2xf32> into tensor<2x4x2xf32>
%221 = arith.constant dense<0.000000e+00> : tensor<2x4x2xf32>
%222 = linalg.batch_matmul ins(%219, %220 : tensor<2x4x4xf32>, tensor<2x4x2xf32>) outs(%221 : tensor<2x4x2xf32>) -> tensor<2x4x2xf32>
%223 = tensor.expand_shape %222 [[0, 1], [2], [3]] : tensor<2x4x2xf32> into tensor<1x2x4x2xf32>
%224 = tensor.empty() : tensor<1x4x2x2xf32>
%225 = linalg.transpose ins(%223 : tensor<1x2x4x2xf32>) outs(%224 : tensor<1x4x2x2xf32>) permutation = [0, 2, 1, 3]
%226 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%227 = tensor.reshape %225(%226) : (tensor<1x4x2x2xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%228 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%229 = linalg.batch_matmul ins(%227, %154 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%228 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%230 = tensor.empty() : tensor<1x4x4xf32>
%231 = linalg.add ins(%143, %229 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%230 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%232 = arith.constant dense<0.000000e+00> : tensor<1x4x8xf32>
%233 = linalg.batch_matmul ins(%231, %156 : tensor<1x4x4xf32>, tensor<1x4x8xf32>) outs(%232 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%236 = arith.constant 0xFF800000 : f32
%234 = tensor.empty() : tensor<1x4x1xf32>
%235 = linalg.fill ins(%236 : f32) outs(%234 : tensor<1x4x1xf32>) -> tensor<1x4x1xf32>
%237 = 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(%233 : tensor<1x4x8xf32>) outs(%235 : tensor<1x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%238 = arith.maximumf %a0, %a1 : f32
linalg.yield %238 : f32
} -> tensor<1x4x1xf32>
%239 = tensor.empty() : tensor<1x4x8xf32>
%240 = 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(%237 : tensor<1x4x1xf32>) outs(%239 : tensor<1x4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x4x8xf32>
%241 = tensor.empty() : tensor<1x4x8xf32>
%242 = linalg.sub ins(%233, %240 : tensor<1x4x8xf32>, tensor<1x4x8xf32>) outs(%241 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%245 = arith.constant 0xFF800000 : f32
%243 = tensor.empty() : tensor<1x4x1xf32>
%244 = linalg.fill ins(%245 : f32) outs(%243 : tensor<1x4x1xf32>) -> tensor<1x4x1xf32>
%246 = 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(%233 : tensor<1x4x8xf32>) outs(%244 : tensor<1x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%247 = arith.maximumf %a0, %a1 : f32
linalg.yield %247 : f32
} -> tensor<1x4x1xf32>
%248 = tensor.empty() : tensor<1x4x8xf32>
%249 = 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(%246 : tensor<1x4x1xf32>) outs(%248 : tensor<1x4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x4x8xf32>
%250 = tensor.empty() : tensor<1x4x8xf32>
%251 = linalg.sub ins(%233, %249 : tensor<1x4x8xf32>, tensor<1x4x8xf32>) outs(%250 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%252 = tensor.empty() : tensor<1x4x8xf32>
%253 = linalg.exp ins(%251 : tensor<1x4x8xf32>) outs(%252 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%256 = arith.constant 0.000000e+00 : f32
%254 = tensor.empty() : tensor<1x4x1xf32>
%255 = linalg.fill ins(%256 : f32) outs(%254 : tensor<1x4x1xf32>) -> tensor<1x4x1xf32>
%257 = 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(%253 : tensor<1x4x8xf32>) outs(%255 : tensor<1x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%258 = arith.addf %a0, %a1 : f32
linalg.yield %258 : f32
} -> tensor<1x4x1xf32>
%259 = tensor.empty() : tensor<1x4x8xf32>
%260 = 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(%257 : tensor<1x4x1xf32>) outs(%259 : tensor<1x4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x4x8xf32>
%261 = tensor.empty() : tensor<1x4x8xf32>
%262 = linalg.log ins(%260 : tensor<1x4x8xf32>) outs(%261 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%263 = tensor.empty() : tensor<1x4x8xf32>
%264 = linalg.sub ins(%242, %262 : tensor<1x4x8xf32>, tensor<1x4x8xf32>) outs(%263 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%265 = arith.constant 0 : i64
%266 = arith.constant 8 : i64
%267 = tensor.empty() : tensor<1x4xi64>
%268 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%126 : tensor<1x4xi64>) outs(%267 : tensor<1x4xi64>) {
^bb0(%candidate: i64, %_: i64):
%269 = arith.cmpi sge, %candidate, %265 : i64
%270 = arith.cmpi slt, %candidate, %266 : i64
%271 = arith.andi %269, %270 : i1
cf.assert %271, "take_along_axis index out of bounds"
linalg.yield %candidate : i64
} -> tensor<1x4xi64>
%272 = tensor.empty() : tensor<1x4xf32>
%273 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%268 : tensor<1x4xi64>) outs(%272 : tensor<1x4xf32>) {
^bb0(%idx: i64, %_: f32):
%274 = linalg.index 0 : index
%275 = linalg.index 1 : index
%276 = arith.index_cast %idx : i64 to index
%277 = tensor.extract %264[%274, %275, %276] : tensor<1x4x8xf32>
linalg.yield %277 : f32
} -> tensor<1x4xf32>
%278 = tensor.empty() : tensor<1x4xf32>
%279 = linalg.negf ins(%273 : tensor<1x4xf32>) outs(%278 : tensor<1x4xf32>) -> tensor<1x4xf32>
return %279 : tensor<1x4xf32>
}
func.func @fn_val_280(%281: tensor<1x4xi64>, %282: tensor<1x4xi64>, %283: tensor<8x4xf32>, %284: tensor<1x1x4x4xf32>, %285: tensor<1x1x1x1xf32>, %286: tensor<4x4xf32>, %287: tensor<4x4xf32>, %288: tensor<4x4xf32>, %289: tensor<4x4xf32>, %290: tensor<4x8xf32>, %291: tensor<1x1xf32>) -> (tensor<1x4xf32>, tensor<8x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x8xf32>) attributes {llvm.emit_c_interface} {
%292 = arith.constant 0 : i64
%293 = arith.constant 8 : i64
%294 = tensor.empty() : tensor<1x4xi64>
%295 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%281 : tensor<1x4xi64>) outs(%294 : tensor<1x4xi64>) {
^bb0(%candidate: i64, %_: i64):
%296 = arith.cmpi sge, %candidate, %292 : i64
%297 = arith.cmpi slt, %candidate, %293 : i64
%298 = arith.andi %296, %297 : i1
cf.assert %298, "gather index out of bounds"
linalg.yield %candidate : i64
} -> tensor<1x4xi64>
%299 = tensor.empty() : tensor<1x4x4xf32>
%300 = 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(%295 : tensor<1x4xi64>) outs(%299 : tensor<1x4x4xf32>) {
^bb0(%idx: i64, %_: f32):
%301 = arith.index_cast %idx : i64 to index
%302 = linalg.index 2 : index
%303 = tensor.extract %283[%301, %302] : tensor<8x4xf32>
linalg.yield %303 : f32
} -> tensor<1x4x4xf32>
%304 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%305 = tensor.reshape %286(%304) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%306 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%307 = tensor.reshape %287(%306) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%308 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%309 = tensor.reshape %288(%308) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%310 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%311 = tensor.reshape %289(%310) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%312 = arith.constant dense<[1, 4, 8]> : tensor<3xi32>
%313 = tensor.reshape %290(%312) : (tensor<4x8xf32>, tensor<3xi32>) -> tensor<1x4x8xf32>
%314 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%315 = linalg.batch_matmul ins(%300, %305 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%314 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%316 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%317 = linalg.batch_matmul ins(%300, %307 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%316 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%318 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%319 = linalg.batch_matmul ins(%300, %309 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%318 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%320 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%321 = tensor.reshape %315(%320) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%322 = tensor.empty() : tensor<1x2x4x2xf32>
%323 = linalg.transpose ins(%321 : tensor<1x4x2x2xf32>) outs(%322 : tensor<1x2x4x2xf32>) permutation = [0, 2, 1, 3]
%324 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%325 = tensor.reshape %317(%324) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%326 = tensor.empty() : tensor<1x2x2x4xf32>
%327 = linalg.transpose ins(%325 : tensor<1x4x2x2xf32>) outs(%326 : tensor<1x2x2x4xf32>) permutation = [0, 2, 3, 1]
%328 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%329 = tensor.reshape %319(%328) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%330 = tensor.empty() : tensor<1x2x4x2xf32>
%331 = linalg.transpose ins(%329 : tensor<1x4x2x2xf32>) outs(%330 : tensor<1x2x4x2xf32>) permutation = [0, 2, 1, 3]
%332 = tensor.collapse_shape %323 [[0, 1], [2], [3]] : tensor<1x2x4x2xf32> into tensor<2x4x2xf32>
%333 = tensor.collapse_shape %327 [[0, 1], [2], [3]] : tensor<1x2x2x4xf32> into tensor<2x2x4xf32>
%334 = arith.constant dense<0.000000e+00> : tensor<2x4x4xf32>
%335 = linalg.batch_matmul ins(%332, %333 : tensor<2x4x2xf32>, tensor<2x2x4xf32>) outs(%334 : tensor<2x4x4xf32>) -> tensor<2x4x4xf32>
%336 = tensor.expand_shape %335 [[0, 1], [2], [3]] : tensor<2x4x4xf32> into tensor<1x2x4x4xf32>
%337 = tensor.empty() : tensor<1x2x4x4xf32>
%338 = 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(%285 : tensor<1x1x1x1xf32>) outs(%337 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%339 = tensor.empty() : tensor<1x2x4x4xf32>
%340 = 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(%284 : tensor<1x1x4x4xf32>) outs(%339 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%341 = tensor.empty() : tensor<1x2x4x4xf32>
%342 = linalg.mul ins(%336, %338 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%341 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%343 = tensor.empty() : tensor<1x2x4x4xf32>
%344 = linalg.add ins(%342, %340 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%343 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%347 = arith.constant 0xFF800000 : f32
%345 = tensor.empty() : tensor<1x2x4x1xf32>
%346 = linalg.fill ins(%347 : f32) outs(%345 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%348 = 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(%344 : tensor<1x2x4x4xf32>) outs(%346 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%349 = arith.maximumf %a0, %a1 : f32
linalg.yield %349 : f32
} -> tensor<1x2x4x1xf32>
%350 = tensor.empty() : tensor<1x2x4x4xf32>
%351 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, d1, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%348 : tensor<1x2x4x1xf32>) outs(%350 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%352 = tensor.empty() : tensor<1x2x4x4xf32>
%353 = linalg.sub ins(%344, %351 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%352 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%354 = tensor.empty() : tensor<1x2x4x4xf32>
%355 = linalg.exp ins(%353 : tensor<1x2x4x4xf32>) outs(%354 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%358 = arith.constant 0xFF800000 : f32
%356 = tensor.empty() : tensor<1x2x4x1xf32>
%357 = linalg.fill ins(%358 : f32) outs(%356 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%359 = 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(%344 : tensor<1x2x4x4xf32>) outs(%357 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%360 = arith.maximumf %a0, %a1 : f32
linalg.yield %360 : f32
} -> tensor<1x2x4x1xf32>
%361 = tensor.empty() : tensor<1x2x4x4xf32>
%362 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, d1, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%359 : tensor<1x2x4x1xf32>) outs(%361 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%363 = tensor.empty() : tensor<1x2x4x4xf32>
%364 = linalg.sub ins(%344, %362 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%363 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%365 = tensor.empty() : tensor<1x2x4x4xf32>
%366 = linalg.exp ins(%364 : tensor<1x2x4x4xf32>) outs(%365 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%369 = arith.constant 0.000000e+00 : f32
%367 = tensor.empty() : tensor<1x2x4x1xf32>
%368 = linalg.fill ins(%369 : f32) outs(%367 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%370 = 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(%366 : tensor<1x2x4x4xf32>) outs(%368 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%371 = arith.addf %a0, %a1 : f32
linalg.yield %371 : f32
} -> tensor<1x2x4x1xf32>
%372 = tensor.empty() : tensor<1x2x4x4xf32>
%373 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, d1, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%370 : tensor<1x2x4x1xf32>) outs(%372 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%374 = tensor.empty() : tensor<1x2x4x4xf32>
%375 = linalg.div ins(%355, %373 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%374 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%376 = tensor.collapse_shape %375 [[0, 1], [2], [3]] : tensor<1x2x4x4xf32> into tensor<2x4x4xf32>
%377 = tensor.collapse_shape %331 [[0, 1], [2], [3]] : tensor<1x2x4x2xf32> into tensor<2x4x2xf32>
%378 = arith.constant dense<0.000000e+00> : tensor<2x4x2xf32>
%379 = linalg.batch_matmul ins(%376, %377 : tensor<2x4x4xf32>, tensor<2x4x2xf32>) outs(%378 : tensor<2x4x2xf32>) -> tensor<2x4x2xf32>
%380 = tensor.expand_shape %379 [[0, 1], [2], [3]] : tensor<2x4x2xf32> into tensor<1x2x4x2xf32>
%381 = tensor.empty() : tensor<1x4x2x2xf32>
%382 = linalg.transpose ins(%380 : tensor<1x2x4x2xf32>) outs(%381 : tensor<1x4x2x2xf32>) permutation = [0, 2, 1, 3]
%383 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%384 = tensor.reshape %382(%383) : (tensor<1x4x2x2xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%385 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%386 = linalg.batch_matmul ins(%384, %311 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%385 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%387 = tensor.empty() : tensor<1x4x4xf32>
%388 = linalg.add ins(%300, %386 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%387 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%389 = arith.constant dense<0.000000e+00> : tensor<1x4x8xf32>
%390 = linalg.batch_matmul ins(%388, %313 : tensor<1x4x4xf32>, tensor<1x4x8xf32>) outs(%389 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%393 = arith.constant 0xFF800000 : f32
%391 = tensor.empty() : tensor<1x4x1xf32>
%392 = linalg.fill ins(%393 : f32) outs(%391 : tensor<1x4x1xf32>) -> tensor<1x4x1xf32>
%394 = 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(%390 : tensor<1x4x8xf32>) outs(%392 : tensor<1x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%395 = arith.maximumf %a0, %a1 : f32
linalg.yield %395 : f32
} -> tensor<1x4x1xf32>
%396 = tensor.empty() : tensor<1x4x8xf32>
%397 = 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(%394 : tensor<1x4x1xf32>) outs(%396 : tensor<1x4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x4x8xf32>
%398 = tensor.empty() : tensor<1x4x8xf32>
%399 = linalg.sub ins(%390, %397 : tensor<1x4x8xf32>, tensor<1x4x8xf32>) outs(%398 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%402 = arith.constant 0xFF800000 : f32
%400 = tensor.empty() : tensor<1x4x1xf32>
%401 = linalg.fill ins(%402 : f32) outs(%400 : tensor<1x4x1xf32>) -> tensor<1x4x1xf32>
%403 = 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(%390 : tensor<1x4x8xf32>) outs(%401 : tensor<1x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%404 = arith.maximumf %a0, %a1 : f32
linalg.yield %404 : f32
} -> tensor<1x4x1xf32>
%405 = tensor.empty() : tensor<1x4x8xf32>
%406 = 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(%403 : tensor<1x4x1xf32>) outs(%405 : tensor<1x4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x4x8xf32>
%407 = tensor.empty() : tensor<1x4x8xf32>
%408 = linalg.sub ins(%390, %406 : tensor<1x4x8xf32>, tensor<1x4x8xf32>) outs(%407 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%409 = tensor.empty() : tensor<1x4x8xf32>
%410 = linalg.exp ins(%408 : tensor<1x4x8xf32>) outs(%409 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%413 = arith.constant 0.000000e+00 : f32
%411 = tensor.empty() : tensor<1x4x1xf32>
%412 = linalg.fill ins(%413 : f32) outs(%411 : tensor<1x4x1xf32>) -> tensor<1x4x1xf32>
%414 = 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(%410 : tensor<1x4x8xf32>) outs(%412 : tensor<1x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%415 = arith.addf %a0, %a1 : f32
linalg.yield %415 : f32
} -> tensor<1x4x1xf32>
%416 = tensor.empty() : tensor<1x4x8xf32>
%417 = 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(%414 : tensor<1x4x1xf32>) outs(%416 : tensor<1x4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x4x8xf32>
%418 = tensor.empty() : tensor<1x4x8xf32>
%419 = linalg.log ins(%417 : tensor<1x4x8xf32>) outs(%418 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%420 = tensor.empty() : tensor<1x4x8xf32>
%421 = linalg.sub ins(%399, %419 : tensor<1x4x8xf32>, tensor<1x4x8xf32>) outs(%420 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%422 = arith.constant 0 : i64
%423 = arith.constant 8 : i64
%424 = tensor.empty() : tensor<1x4xi64>
%425 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%282 : tensor<1x4xi64>) outs(%424 : tensor<1x4xi64>) {
^bb0(%candidate: i64, %_: i64):
%426 = arith.cmpi sge, %candidate, %422 : i64
%427 = arith.cmpi slt, %candidate, %423 : i64
%428 = arith.andi %426, %427 : i1
cf.assert %428, "take_along_axis index out of bounds"
linalg.yield %candidate : i64
} -> tensor<1x4xi64>
%429 = tensor.empty() : tensor<1x4xf32>
%430 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%425 : tensor<1x4xi64>) outs(%429 : tensor<1x4xf32>) {
^bb0(%idx: i64, %_: f32):
%431 = linalg.index 0 : index
%432 = linalg.index 1 : index
%433 = arith.index_cast %idx : i64 to index
%434 = tensor.extract %421[%431, %432, %433] : tensor<1x4x8xf32>
linalg.yield %434 : f32
} -> tensor<1x4xf32>
%435 = tensor.empty() : tensor<1x4xf32>
%436 = linalg.negf ins(%430 : tensor<1x4xf32>) outs(%435 : tensor<1x4xf32>) -> tensor<1x4xf32>
%437 = arith.constant dense<1.000000e+00> : tensor<1x4xf32>
%438 = arith.constant dense<1.000000e+00> : tensor<1x4xf32>
%439 = tensor.empty() : tensor<1x4xf32>
%440 = linalg.negf ins(%438 : tensor<1x4xf32>) outs(%439 : tensor<1x4xf32>) -> tensor<1x4xf32>
%441 = arith.constant 0 : i64
%442 = arith.constant 8 : i64
%443 = tensor.empty() : tensor<1x4xi64>
%444 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%282 : tensor<1x4xi64>) outs(%443 : tensor<1x4xi64>) {
^bb0(%candidate: i64, %_: i64):
%445 = arith.cmpi sge, %candidate, %441 : i64
%446 = arith.cmpi slt, %candidate, %442 : i64
%447 = arith.andi %445, %446 : i1
cf.assert %447, "scatter_along_axis index out of bounds"
linalg.yield %candidate : i64
} -> tensor<1x4xi64>
%448 = tensor.empty() : tensor<1x4x8xf32>
%449 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%444, %440 : tensor<1x4xi64>, tensor<1x4xf32>) outs(%448 : tensor<1x4x8xf32>) {
^bb0(%idx: i64, %upd: f32, %_: f32):
%450 = linalg.index 2 : index
%451 = arith.index_cast %450 : index to i64
%452 = arith.cmpi eq, %idx, %451 : i64
%453 = arith.constant 0.000000e+00 : f32
%454 = arith.select %452, %upd, %453 : f32
linalg.yield %454 : f32
} -> tensor<1x4x8xf32>
%455 = tensor.empty() : tensor<1x4x8xf32>
%456 = linalg.exp ins(%421 : tensor<1x4x8xf32>) outs(%455 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%459 = arith.constant 0.000000e+00 : f32
%457 = tensor.empty() : tensor<1x4x1xf32>
%458 = linalg.fill ins(%459 : f32) outs(%457 : tensor<1x4x1xf32>) -> tensor<1x4x1xf32>
%460 = 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(%449 : tensor<1x4x8xf32>) outs(%458 : tensor<1x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%461 = arith.addf %a0, %a1 : f32
linalg.yield %461 : f32
} -> tensor<1x4x1xf32>
%462 = tensor.empty() : tensor<1x4x8xf32>
%463 = 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(%460 : tensor<1x4x1xf32>) outs(%462 : tensor<1x4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x4x8xf32>
%464 = tensor.empty() : tensor<1x4x8xf32>
%465 = linalg.mul ins(%463, %456 : tensor<1x4x8xf32>, tensor<1x4x8xf32>) outs(%464 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%466 = tensor.empty() : tensor<1x4x8xf32>
%467 = linalg.sub ins(%449, %465 : tensor<1x4x8xf32>, tensor<1x4x8xf32>) outs(%466 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%468 = tensor.empty() : tensor<1x4x4xf32>
%469 = linalg.transpose ins(%388 : tensor<1x4x4xf32>) outs(%468 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%470 = arith.constant dense<0.000000e+00> : tensor<1x4x8xf32>
%471 = linalg.batch_matmul ins(%469, %467 : tensor<1x4x4xf32>, tensor<1x4x8xf32>) outs(%470 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%472 = tensor.empty() : tensor<1x8x4xf32>
%473 = linalg.transpose ins(%313 : tensor<1x4x8xf32>) outs(%472 : tensor<1x8x4xf32>) permutation = [0, 2, 1]
%474 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%475 = linalg.batch_matmul ins(%467, %473 : tensor<1x4x8xf32>, tensor<1x8x4xf32>) outs(%474 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%476 = tensor.empty() : tensor<1x4x4xf32>
%477 = linalg.transpose ins(%384 : tensor<1x4x4xf32>) outs(%476 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%478 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%479 = linalg.batch_matmul ins(%477, %475 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%478 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%480 = tensor.empty() : tensor<1x4x4xf32>
%481 = linalg.transpose ins(%311 : tensor<1x4x4xf32>) outs(%480 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%482 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%483 = linalg.batch_matmul ins(%475, %481 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%482 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%484 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%485 = tensor.reshape %483(%484) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%486 = tensor.empty() : tensor<1x2x4x2xf32>
%487 = linalg.transpose ins(%485 : tensor<1x4x2x2xf32>) outs(%486 : tensor<1x2x4x2xf32>) permutation = [0, 2, 1, 3]
%488 = tensor.empty() : tensor<1x2x4x4xf32>
%489 = linalg.transpose ins(%375 : tensor<1x2x4x4xf32>) outs(%488 : tensor<1x2x4x4xf32>) permutation = [0, 1, 3, 2]
%490 = tensor.collapse_shape %489 [[0, 1], [2], [3]] : tensor<1x2x4x4xf32> into tensor<2x4x4xf32>
%491 = tensor.collapse_shape %487 [[0, 1], [2], [3]] : tensor<1x2x4x2xf32> into tensor<2x4x2xf32>
%492 = arith.constant dense<0.000000e+00> : tensor<2x4x2xf32>
%493 = linalg.batch_matmul ins(%490, %491 : tensor<2x4x4xf32>, tensor<2x4x2xf32>) outs(%492 : tensor<2x4x2xf32>) -> tensor<2x4x2xf32>
%494 = tensor.expand_shape %493 [[0, 1], [2], [3]] : tensor<2x4x2xf32> into tensor<1x2x4x2xf32>
%495 = tensor.empty() : tensor<1x2x2x4xf32>
%496 = linalg.transpose ins(%331 : tensor<1x2x4x2xf32>) outs(%495 : tensor<1x2x2x4xf32>) permutation = [0, 1, 3, 2]
%497 = tensor.collapse_shape %487 [[0, 1], [2], [3]] : tensor<1x2x4x2xf32> into tensor<2x4x2xf32>
%498 = tensor.collapse_shape %496 [[0, 1], [2], [3]] : tensor<1x2x2x4xf32> into tensor<2x2x4xf32>
%499 = arith.constant dense<0.000000e+00> : tensor<2x4x4xf32>
%500 = linalg.batch_matmul ins(%497, %498 : tensor<2x4x2xf32>, tensor<2x2x4xf32>) outs(%499 : tensor<2x4x4xf32>) -> tensor<2x4x4xf32>
%501 = tensor.expand_shape %500 [[0, 1], [2], [3]] : tensor<2x4x4xf32> into tensor<1x2x4x4xf32>
%502 = tensor.empty() : tensor<1x2x4x4xf32>
%503 = linalg.mul ins(%375, %501 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%502 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%506 = arith.constant 0.000000e+00 : f32
%504 = tensor.empty() : tensor<1x2x4x1xf32>
%505 = linalg.fill ins(%506 : f32) outs(%504 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%507 = 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(%503 : tensor<1x2x4x4xf32>) outs(%505 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%508 = arith.addf %a0, %a1 : f32
linalg.yield %508 : f32
} -> tensor<1x2x4x1xf32>
%509 = tensor.empty() : tensor<1x2x4x4xf32>
%510 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, d1, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%507 : tensor<1x2x4x1xf32>) outs(%509 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%511 = tensor.empty() : tensor<1x2x4x4xf32>
%512 = linalg.sub ins(%501, %510 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%511 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%513 = tensor.empty() : tensor<1x2x4x4xf32>
%514 = linalg.mul ins(%375, %512 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%513 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%515 = tensor.empty() : tensor<1x2x4x4xf32>
%516 = linalg.mul ins(%514, %338 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%515 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%517 = tensor.empty() : tensor<1x2x4x4xf32>
%518 = linalg.mul ins(%514, %336 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%517 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%521 = arith.constant 0.000000e+00 : f32
%519 = tensor.empty() : tensor<1x1x4x4xf32>
%520 = linalg.fill ins(%521 : f32) outs(%519 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%522 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, 0, d2, d3)>],
iterator_types = ["parallel", "reduction", "parallel", "parallel"]
} ins(%514 : tensor<1x2x4x4xf32>) outs(%520 : tensor<1x1x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
%523 = arith.addf %a0, %a1 : f32
linalg.yield %523 : f32
} -> tensor<1x1x4x4xf32>
%526 = arith.constant 0.000000e+00 : f32
%524 = tensor.empty() : tensor<1x2x4x1xf32>
%525 = linalg.fill ins(%526 : f32) outs(%524 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%527 = 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(%518 : tensor<1x2x4x4xf32>) outs(%525 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%528 = arith.addf %a0, %a1 : f32
linalg.yield %528 : f32
} -> tensor<1x2x4x1xf32>
%531 = arith.constant 0.000000e+00 : f32
%529 = tensor.empty() : tensor<1x2x1x1xf32>
%530 = linalg.fill ins(%531 : f32) outs(%529 : tensor<1x2x1x1xf32>) -> tensor<1x2x1x1xf32>
%532 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, 0, d3)>],
iterator_types = ["parallel", "parallel", "reduction", "parallel"]
} ins(%527 : tensor<1x2x4x1xf32>) outs(%530 : tensor<1x2x1x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%533 = arith.addf %a0, %a1 : f32
linalg.yield %533 : f32
} -> tensor<1x2x1x1xf32>
%536 = arith.constant 0.000000e+00 : f32
%534 = tensor.empty() : tensor<1x1x1x1xf32>
%535 = linalg.fill ins(%536 : f32) outs(%534 : tensor<1x1x1x1xf32>) -> tensor<1x1x1x1xf32>
%537 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, 0, d2, d3)>],
iterator_types = ["parallel", "reduction", "parallel", "parallel"]
} ins(%532 : tensor<1x2x1x1xf32>) outs(%535 : tensor<1x1x1x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%538 = arith.addf %a0, %a1 : f32
linalg.yield %538 : f32
} -> tensor<1x1x1x1xf32>
%539 = tensor.empty() : tensor<1x2x2x4xf32>
%540 = linalg.transpose ins(%323 : tensor<1x2x4x2xf32>) outs(%539 : tensor<1x2x2x4xf32>) permutation = [0, 1, 3, 2]
%541 = tensor.collapse_shape %540 [[0, 1], [2], [3]] : tensor<1x2x2x4xf32> into tensor<2x2x4xf32>
%542 = tensor.collapse_shape %516 [[0, 1], [2], [3]] : tensor<1x2x4x4xf32> into tensor<2x4x4xf32>
%543 = arith.constant dense<0.000000e+00> : tensor<2x2x4xf32>
%544 = linalg.batch_matmul ins(%541, %542 : tensor<2x2x4xf32>, tensor<2x4x4xf32>) outs(%543 : tensor<2x2x4xf32>) -> tensor<2x2x4xf32>
%545 = tensor.expand_shape %544 [[0, 1], [2], [3]] : tensor<2x2x4xf32> into tensor<1x2x2x4xf32>
%546 = tensor.empty() : tensor<1x2x4x2xf32>
%547 = linalg.transpose ins(%327 : tensor<1x2x2x4xf32>) outs(%546 : tensor<1x2x4x2xf32>) permutation = [0, 1, 3, 2]
%548 = tensor.collapse_shape %516 [[0, 1], [2], [3]] : tensor<1x2x4x4xf32> into tensor<2x4x4xf32>
%549 = tensor.collapse_shape %547 [[0, 1], [2], [3]] : tensor<1x2x4x2xf32> into tensor<2x4x2xf32>
%550 = arith.constant dense<0.000000e+00> : tensor<2x4x2xf32>
%551 = linalg.batch_matmul ins(%548, %549 : tensor<2x4x4xf32>, tensor<2x4x2xf32>) outs(%550 : tensor<2x4x2xf32>) -> tensor<2x4x2xf32>
%552 = tensor.expand_shape %551 [[0, 1], [2], [3]] : tensor<2x4x2xf32> into tensor<1x2x4x2xf32>
%553 = tensor.empty() : tensor<1x4x2x2xf32>
%554 = linalg.transpose ins(%494 : tensor<1x2x4x2xf32>) outs(%553 : tensor<1x4x2x2xf32>) permutation = [0, 2, 1, 3]
%555 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%556 = tensor.reshape %554(%555) : (tensor<1x4x2x2xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%557 = tensor.empty() : tensor<1x4x2x2xf32>
%558 = linalg.transpose ins(%545 : tensor<1x2x2x4xf32>) outs(%557 : tensor<1x4x2x2xf32>) permutation = [0, 3, 1, 2]
%559 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%560 = tensor.reshape %558(%559) : (tensor<1x4x2x2xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%561 = tensor.empty() : tensor<1x4x2x2xf32>
%562 = linalg.transpose ins(%552 : tensor<1x2x4x2xf32>) outs(%561 : tensor<1x4x2x2xf32>) permutation = [0, 2, 1, 3]
%563 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%564 = tensor.reshape %562(%563) : (tensor<1x4x2x2xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%565 = tensor.empty() : tensor<1x4x4xf32>
%566 = linalg.transpose ins(%300 : tensor<1x4x4xf32>) outs(%565 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%567 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%568 = linalg.batch_matmul ins(%566, %556 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%567 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%569 = tensor.empty() : tensor<1x4x4xf32>
%570 = linalg.transpose ins(%309 : tensor<1x4x4xf32>) outs(%569 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%571 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%572 = linalg.batch_matmul ins(%556, %570 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%571 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%573 = tensor.empty() : tensor<1x4x4xf32>
%574 = linalg.transpose ins(%300 : tensor<1x4x4xf32>) outs(%573 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%575 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%576 = linalg.batch_matmul ins(%574, %560 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%575 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%577 = tensor.empty() : tensor<1x4x4xf32>
%578 = linalg.transpose ins(%307 : tensor<1x4x4xf32>) outs(%577 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%579 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%580 = linalg.batch_matmul ins(%560, %578 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%579 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%581 = tensor.empty() : tensor<1x4x4xf32>
%582 = linalg.transpose ins(%300 : tensor<1x4x4xf32>) outs(%581 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%583 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%584 = linalg.batch_matmul ins(%582, %564 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%583 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%585 = tensor.empty() : tensor<1x4x4xf32>
%586 = linalg.transpose ins(%305 : tensor<1x4x4xf32>) outs(%585 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%587 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%588 = linalg.batch_matmul ins(%564, %586 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%587 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%589 = arith.constant dense<[4, 8]> : tensor<2xi32>
%590 = tensor.reshape %471(%589) : (tensor<1x4x8xf32>, tensor<2xi32>) -> tensor<4x8xf32>
%591 = arith.constant dense<[4, 4]> : tensor<2xi32>
%592 = tensor.reshape %479(%591) : (tensor<1x4x4xf32>, tensor<2xi32>) -> tensor<4x4xf32>
%593 = arith.constant dense<[4, 4]> : tensor<2xi32>
%594 = tensor.reshape %568(%593) : (tensor<1x4x4xf32>, tensor<2xi32>) -> tensor<4x4xf32>
%595 = arith.constant dense<[4, 4]> : tensor<2xi32>
%596 = tensor.reshape %576(%595) : (tensor<1x4x4xf32>, tensor<2xi32>) -> tensor<4x4xf32>
%597 = arith.constant dense<[4, 4]> : tensor<2xi32>
%598 = tensor.reshape %584(%597) : (tensor<1x4x4xf32>, tensor<2xi32>) -> tensor<4x4xf32>
%599 = tensor.empty() : tensor<1x4x4xf32>
%600 = linalg.add ins(%475, %572 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%599 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%601 = tensor.empty() : tensor<1x4x4xf32>
%602 = linalg.add ins(%600, %580 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%601 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%603 = tensor.empty() : tensor<1x4x4xf32>
%604 = linalg.add ins(%602, %588 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%603 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%605 = arith.constant 0 : i64
%606 = arith.constant 8 : i64
%607 = tensor.empty() : tensor<1x4xi64>
%608 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%281 : tensor<1x4xi64>) outs(%607 : tensor<1x4xi64>) {
^bb0(%candidate: i64, %_: i64):
%609 = arith.cmpi sge, %candidate, %605 : i64
%610 = arith.cmpi slt, %candidate, %606 : i64
%611 = arith.andi %609, %610 : i1
cf.assert %611, "scatter_add index out of bounds"
linalg.yield %candidate : i64
} -> tensor<1x4xi64>
%612 = tensor.collapse_shape %608 [[0, 1]] : tensor<1x4xi64> into tensor<4xi64>
%613 = tensor.collapse_shape %604 [[0, 1], [2]] : tensor<1x4x4xf32> into tensor<4x4xf32>
%614 = arith.constant dense<0.000000e+00> : tensor<8x4xf32>
%615 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d2)>, affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%612, %613 : tensor<4xi64>, tensor<4x4xf32>) outs(%614 : tensor<8x4xf32>) {
^bb0(%idx: i64, %upd: f32, %acc: f32):
%616 = linalg.index 0 : index
%617 = arith.index_cast %616 : index to i64
%618 = arith.cmpi eq, %idx, %617 : i64
%619 = arith.constant 0.000000e+00 : f32
%620 = arith.select %618, %upd, %619 : f32
%621 = arith.addf %acc, %620 : f32
linalg.yield %621 : f32
} -> tensor<8x4xf32>
%622 = tensor.empty() : tensor<8x4xf32>
%623 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%291 : tensor<1x1xf32>) outs(%622 : tensor<8x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<8x4xf32>
%624 = tensor.empty() : tensor<8x4xf32>
%625 = linalg.mul ins(%623, %615 : tensor<8x4xf32>, tensor<8x4xf32>) outs(%624 : tensor<8x4xf32>) -> tensor<8x4xf32>
%626 = tensor.empty() : tensor<8x4xf32>
%627 = linalg.sub ins(%283, %625 : tensor<8x4xf32>, tensor<8x4xf32>) outs(%626 : tensor<8x4xf32>) -> tensor<8x4xf32>
%628 = tensor.empty() : tensor<4x4xf32>
%629 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%291 : tensor<1x1xf32>) outs(%628 : tensor<4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x4xf32>
%630 = tensor.empty() : tensor<4x4xf32>
%631 = linalg.mul ins(%629, %598 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%630 : tensor<4x4xf32>) -> tensor<4x4xf32>
%632 = tensor.empty() : tensor<4x4xf32>
%633 = linalg.sub ins(%286, %631 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%632 : tensor<4x4xf32>) -> tensor<4x4xf32>
%634 = tensor.empty() : tensor<4x4xf32>
%635 = linalg.mul ins(%629, %596 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%634 : tensor<4x4xf32>) -> tensor<4x4xf32>
%636 = tensor.empty() : tensor<4x4xf32>
%637 = linalg.sub ins(%287, %635 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%636 : tensor<4x4xf32>) -> tensor<4x4xf32>
%638 = tensor.empty() : tensor<4x4xf32>
%639 = linalg.mul ins(%629, %594 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%638 : tensor<4x4xf32>) -> tensor<4x4xf32>
%640 = tensor.empty() : tensor<4x4xf32>
%641 = linalg.sub ins(%288, %639 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%640 : tensor<4x4xf32>) -> tensor<4x4xf32>
%642 = tensor.empty() : tensor<4x4xf32>
%643 = linalg.mul ins(%629, %592 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%642 : tensor<4x4xf32>) -> tensor<4x4xf32>
%644 = tensor.empty() : tensor<4x4xf32>
%645 = linalg.sub ins(%289, %643 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%644 : tensor<4x4xf32>) -> tensor<4x4xf32>
%646 = tensor.empty() : tensor<4x8xf32>
%647 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%291 : tensor<1x1xf32>) outs(%646 : tensor<4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x8xf32>
%648 = tensor.empty() : tensor<4x8xf32>
%649 = linalg.mul ins(%647, %590 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%648 : tensor<4x8xf32>) -> tensor<4x8xf32>
%650 = tensor.empty() : tensor<4x8xf32>
%651 = linalg.sub ins(%290, %649 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%650 : tensor<4x8xf32>) -> tensor<4x8xf32>
return %436, %627, %633, %637, %641, %645, %651 : tensor<1x4xf32>, tensor<8x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x8xf32>
}
func.func @main() -> (tensor<1x4xf32>, tensor<8x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x8xf32>) {
%652 = arith.constant dense<0> : tensor<1x4xi64>
%653 = arith.constant dense<0> : tensor<1x4xi64>
%654 = arith.constant dense<0.000000e+00> : tensor<8x4xf32>
%655 = arith.constant dense<0.000000e+00> : tensor<1x1x4x4xf32>
%656 = arith.constant dense<0.000000e+00> : tensor<1x1x1x1xf32>
%657 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%658 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%659 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%660 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%661 = arith.constant dense<0.000000e+00> : tensor<4x8xf32>
%662 = arith.constant dense<0.000000e+00> : tensor<1x1xf32>
%663, %664, %665, %666, %667, %668, %669 = func.call @fn_val_280(%652, %653, %654, %655, %656, %657, %658, %659, %660, %661, %662) : (tensor<1x4xi64>, tensor<1x4xi64>, tensor<8x4xf32>, tensor<1x1x4x4xf32>, tensor<1x1x1x1xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x8xf32>, tensor<1x1xf32>) -> (tensor<1x4xf32>, tensor<8x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x8xf32>)
return %663, %664, %665, %666, %667, %668, %669 : tensor<1x4xf32>, tensor<8x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x8xf32>
}
}