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